Spaces:
Sleeping
Sleeping
Commit
•
973f2c1
1
Parent(s):
5236358
Update tinysam/build_sam.py
Browse files- tinysam/build_sam.py +2 -1
tinysam/build_sam.py
CHANGED
@@ -88,7 +88,8 @@ def build_sam_vit_t(checkpoint=None):
|
|
88 |
mobile_sam.eval()
|
89 |
if checkpoint is not None:
|
90 |
with open(checkpoint, "rb") as f:
|
91 |
-
|
|
|
92 |
mobile_sam.load_state_dict(state_dict)
|
93 |
return mobile_sam
|
94 |
|
|
|
88 |
mobile_sam.eval()
|
89 |
if checkpoint is not None:
|
90 |
with open(checkpoint, "rb") as f:
|
91 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
92 |
+
state_dict = torch.load(f, map_location=device)
|
93 |
mobile_sam.load_state_dict(state_dict)
|
94 |
return mobile_sam
|
95 |
|