Spaces:
Running
Running
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)
|