Hugging Face's logo

Spaces:
Yakova
/
Mobile-SAM
Running

App Files Community
YouLiXiya commited on
Commit
4508aa9
•
1 Parent(s): 6516d18

Update segment_anything/segment_anything/mobile_encoder/setup_mobile_sam.py

Browse files
Files changed (1) hide show
  1. segment_anything/segment_anything/mobile_encoder/setup_mobile_sam.py +1 -1
segment_anything/segment_anything/mobile_encoder/setup_mobile_sam.py CHANGED
@@ -44,7 +44,7 @@ def setup_model():
44
  return mobile_sam
45
  def load_mobile_sam(mobile_sam_checkpoint_path,
46
  device='cuda'):
47
- checkpoint = torch.load(mobile_sam_checkpoint_path)
48
  mobile_sam = setup_model()
49
  mobile_sam.load_state_dict(checkpoint, strict=True)
50
  mobile_sam.to(device=device)
 
44
  return mobile_sam
45
  def load_mobile_sam(mobile_sam_checkpoint_path,
46
  device='cuda'):
47
+ checkpoint = torch.load(mobile_sam_checkpoint_path, map_location='cpu')
48
  mobile_sam = setup_model()
49
  mobile_sam.load_state_dict(checkpoint, strict=True)
50
  mobile_sam.to(device=device)