fix issue with version splitting (#1074)

MouseLand · Feb 7, 2025 · 3d95187 · 3d95187
1 parent 81c3bfb
commit 3d95187
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cellpose/dynamics.py b/cellpose/dynamics.py
@@ -620,7 +620,7 @@ def remove_bad_flow_masks(masks, flows, threshold=0.4, device=torch.device("cpu"
     device0 = device
     if masks.size > 10000 * 10000 and (device is not None and device.type == "cuda"):
 
-        major_version, minor_version, _ = torch.__version__.split(".")
+        major_version, minor_version = torch.__version__.split(".")[:2]
         torch.cuda.empty_cache()
         if major_version == "1" and int(minor_version) < 10:
             # for PyTorch version lower than 1.10