Merge pull request #163 from JoOkuma/device-fix

Using predictor device when setting image embedding
computational-cell-analytics · Aug 24, 2023 · bf82e31 · bf82e31
2 parents cad7c8f + 8a72e87
commit bf82e31
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/micro_sam/util.py b/micro_sam/util.py
@@ -560,7 +560,7 @@ def set_precomputed(
         i: Index for the image data. Required if `image` has three spatial dimensions
             or a time dimension and two spatial dimensions.
     """
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = predictor.device
     features = image_embeddings["features"]
 
     assert features.ndim in (4, 5)