fixing setup.py removal of pytorch for new torch 2.0 release

MouseLand · Feb 18, 2023 · e559dae · e559dae
1 parent e86cd01
commit e559dae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -34,8 +34,8 @@
 try:
     import torch
     a = torch.ones(2, 3)
-    version = int(torch.__version__.split(".")[1])
-    if version >= 6:
+    major_version, minor_version, _ = torch.__version__.split(".")
+    if major_version == "2" or int(minor_version) >= 6:
         install_deps.remove("torch>=1.6")
 except:
     pass