diff --git a/py/sam_hq/build_sam_hq.py b/py/sam_hq/build_sam_hq.py
index c8f41f2..a4608b8 100644
--- a/py/sam_hq/build_sam_hq.py
+++ b/py/sam_hq/build_sam_hq.py
@@ -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():