Merge pull request #1 from weigertlab/mps_fix

Mps fix
weigertlab · Aug 29, 2024 · 44d6bfd · 44d6bfd
2 parents b044e7b + cfb26b6
commit 44d6bfd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/napari_spotiflow/_dock_widget.py b/napari_spotiflow/_dock_widget.py
@@ -186,8 +186,9 @@ def plugin (
     ) -> list[napari.types.LayerDataTuple]:
         if image_axes == "":
             raise RuntimeError("Invalid axes order. If your input is 2D, please set the 2D mode. If your input is 3D, please set the 3D mode.")
-        should_use_mps = torch.backends.mps.is_available() and (not IS_3D or (os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") is not None and os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") != "0"))
-        DEVICE_STR = "cuda" if torch.cuda.is_available() else "mps" if should_use_mps and os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") else "cpu"
+        should_use_mps = torch.backends.mps.is_available() and not IS_3D and not os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") == "0"
+        DEVICE_STR = "cuda" if torch.cuda.is_available() else "mps" if should_use_mps else "cpu"
+        print(f'using device {DEVICE_STR}')
 
         model = get_model(
             model_type,