From 1891fedd49c647d8e2379422d5b790a45bcd0951 Mon Sep 17 00:00:00 2001
From: Logan Adams <loadams@microsoft.com>
Date: Fri, 27 Sep 2024 09:44:57 -0700
Subject: [PATCH 1/3] Fix MLU op builder torch install

---
 op_builder/mlu/fused_adam.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/op_builder/mlu/fused_adam.py b/op_builder/mlu/fused_adam.py
index 0198db7cb276..76b94f949481 100644
--- a/op_builder/mlu/fused_adam.py
+++ b/op_builder/mlu/fused_adam.py
@@ -13,6 +13,12 @@ class MLUFusedAdam:
     @staticmethod
     def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode,
                           bias_correction, weight_decay, *args):
+        try:
+            import torch
+        except ImportError:
+            if verbose:
+                self.warning("Please install torch if trying to pre-compile GDS")
+            return False
 
         torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2],
                                        tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction,

From 5ea5ab08d3f7913e00fc102a7a7e01bd7d3ec4d2 Mon Sep 17 00:00:00 2001
From: Logan Adams <loadams@microsoft.com>
Date: Fri, 27 Sep 2024 09:56:17 -0700
Subject: [PATCH 2/3] Formatting/syntax error

---
 op_builder/mlu/fused_adam.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/op_builder/mlu/fused_adam.py b/op_builder/mlu/fused_adam.py
index 76b94f949481..0751f3f14b60 100644
--- a/op_builder/mlu/fused_adam.py
+++ b/op_builder/mlu/fused_adam.py
@@ -5,7 +5,11 @@
 # DeepSpeed Team
 
 from .builder import MLUOpBuilder
-import torch
+
+try:
+    import torch
+except: ImportError as e:
+    pass
 
 
 class MLUFusedAdam:
@@ -13,12 +17,6 @@ class MLUFusedAdam:
     @staticmethod
     def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode,
                           bias_correction, weight_decay, *args):
-        try:
-            import torch
-        except ImportError:
-            if verbose:
-                self.warning("Please install torch if trying to pre-compile GDS")
-            return False
 
         torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2],
                                        tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction,

From 54ce016dbc500dfa624c6e69e7eaefd1e33f50dc Mon Sep 17 00:00:00 2001
From: Logan Adams <loadams@microsoft.com>
Date: Fri, 27 Sep 2024 09:57:17 -0700
Subject: [PATCH 3/3] Fix syntax

---
 op_builder/mlu/fused_adam.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/op_builder/mlu/fused_adam.py b/op_builder/mlu/fused_adam.py
index 0751f3f14b60..2343e55568c4 100644
--- a/op_builder/mlu/fused_adam.py
+++ b/op_builder/mlu/fused_adam.py
@@ -8,7 +8,7 @@
 
 try:
     import torch
-except: ImportError as e:
+except ImportError as e:
     pass