move setup.py import

pinellolab · Oct 15, 2021 · 7cf6c4f · 7cf6c4f
1 parent 029eec5
commit 7cf6c4f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -9,13 +9,13 @@
 import os
 
 from setuptools import setup
-from torch.utils import cpp_extension
 
 
 if __name__ == "__main__":
     if int(os.getenv("PBG_INSTALL_CPP", 0)) == 0:
         setup()
     else:
+        from torch.utils import cpp_extension
         setup(
             ext_modules=[
                 cpp_extension.CppExtension(