move torch import (#4468)

microsoft · Oct 9, 2023 · 6b634d0 · 6b634d0
1 parent 59d62d6
commit 6b634d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/op_builder/async_io.py b/op_builder/async_io.py
@@ -5,7 +5,6 @@
 
 import distutils.spawn
 import subprocess
-import torch
 
 from .builder import OpBuilder
 
@@ -36,6 +35,7 @@ def cxx_args(self):
         # -O0 for improved debugging, since performance is bound by I/O
         CPU_ARCH = self.cpu_arch()
         SIMD_WIDTH = self.simd_width()
+        import torch  # Keep this import here to avoid errors when building DeepSpeed wheel without torch installed
         TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2])
         if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1:
             CPP_STD = '-std=c++17'