YouLiXiya commited on
Commit
4508aa9
1 Parent(s): 6516d18

Update segment_anything/segment_anything/mobile_encoder/setup_mobile_sam.py

Browse files
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)