change map location to CPU

inference_sam.py

@@ -33,11 +33,22 @@ if not os.path.exists('model'):
      print("warning! A read token in env variables is needed for authentication.")
   snapshot_download(repo_id=REPO_ID, token=token,repo_type='model',local_dir='model')
 
+
+original_torch_load = torch.load
+
+def patched_torch_load(*args, **kwargs):
+    kwargs['map_location'] = device
+    return original_torch_load(*args, **kwargs)
+
+torch.load = patched_torch_load
+
 model_path = os.path.join('model', 'sam_02-06_dice_mse_0.pth')
-sam = sam_model_registry["default"](model_path,map_location=device)
+sam = sam_model_registry["default"](model_path)
 sam.to(device) #sam.cuda()
 predictor = SamPredictor(sam)
 
+torch.load = original_torch_load
+
 
 from torch.nn import functional as F