andy-wyx commited on
Commit
f6436fb
·
1 Parent(s): 70c5a8b

change map location to CPU

Browse files
Files changed (1) hide show
  1. inference_sam.py +1 -1
inference_sam.py CHANGED
@@ -34,7 +34,7 @@ if not os.path.exists('model'):
34
  snapshot_download(repo_id=REPO_ID, token=token,repo_type='model',local_dir='model')
35
 
36
  model_path = os.path.join('model', 'sam_02-06_dice_mse_0.pth')
37
- sam = sam_model_registry["default"](model_path)
38
  sam.to(device) #sam.cuda()
39
  predictor = SamPredictor(sam)
40
 
 
34
  snapshot_download(repo_id=REPO_ID, token=token,repo_type='model',local_dir='model')
35
 
36
  model_path = os.path.join('model', 'sam_02-06_dice_mse_0.pth')
37
+ sam = sam_model_registry["default"](model_path,map_location=device)
38
  sam.to(device) #sam.cuda()
39
  predictor = SamPredictor(sam)
40