Update tinysam/build_sam.py

tinysam/build_sam.py

@@ -88,7 +88,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)
         mobile_sam.load_state_dict(state_dict)
     return mobile_sam