using predictor device when setting image embedding

computational-cell-analytics · Aug 24, 2023 · 8a72e87 · 8a72e87
1 parent cad7c8f
commit 8a72e87
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)