Merge pull request #21 from coderobe/main

load sam_hq to cpu if cuda is unavailable
2025-01-05 14:56:21 +08:00
parent b0336283f7 1b8f5183c0
commit 4991451b73
1 changed files with 4 additions and 1 deletions
@@ -64,7 +64,10 @@ 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)
+            if torch.cuda.is_available():
+                state_dict = torch.load(f)
+            else:
+                state_dict = torch.load(f, map_location=torch.device('cpu'))
        info = sam.load_state_dict(state_dict, strict=False)
        print(info)
    for _, p in sam.named_parameters():