From 8a72e87bc02ea2a786a5bb2a5c85818356a957dc Mon Sep 17 00:00:00 2001 From: Jordao Bragantini Date: Thu, 24 Aug 2023 11:17:51 -0700 Subject: [PATCH] using predictor device when setting image embedding --- micro_sam/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/micro_sam/util.py b/micro_sam/util.py index 1f705e94..1ea21303 100644 --- 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)