fix sam hq and mobile sam loading with non cuda system (#190)

continue-revolution · Feb 7, 2024 · b5be424 · b5be424
1 parent 2ba9dde
commit b5be424
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
 models/*/*.pth
+models/*/*.pt
 __pycache__
-annotator
+annotator
diff --git a/sam_hq/build_sam_hq.py b/sam_hq/build_sam_hq.py
@@ -64,7 +64,7 @@ def _load_sam_checkpoint(sam: Sam, checkpoint=None):
     sam.eval()
     if checkpoint is not None:
         with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
+            state_dict = torch.load(f, map_location="cpu")
         info = sam.load_state_dict(state_dict, strict=False)
         print(info)
     for _, p in sam.named_parameters():