Merge pull request #706 from mrariden/gpu_check_flow_masks

make GPU check fix
MouseLand · May 14, 2023 · ca88d6a · ca88d6a
2 parents b709064 + de61b8b
commit ca88d6a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cellpose/dynamics.py b/cellpose/dynamics.py
@@ -577,7 +577,7 @@ def remove_bad_flow_masks(masks, flows, threshold=0.4, use_gpu=False, device=Non
         size [Ly x Lx] or [Lz x Ly x Lx]
     
     """
-    if masks.size > 10000*10000:
+    if masks.size > 10000*10000 and use_gpu:
 
         major_version, minor_version, _ = torch.__version__.split(".")