cpu

Author: lkeab

segment_anything/build_sam.py

@@ -89,7 +89,8 @@ def build_sam_vit_t(checkpoint=None):
     mobile_sam.eval()
     if checkpoint is not None:
         with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            state_dict = torch.load(f, map_location=device)
         info = mobile_sam.load_state_dict(state_dict, strict=False)
         print(info)
     for n, p in mobile_sam.named_parameters():