Spaces:
Running
on
Zero
Running
on
Zero
Commit
·
12e7eb6
1
Parent(s):
6192832
fix load state
Browse files
XPart/partgen/bbox_estimator/auto_mask_api.py
CHANGED
|
@@ -1363,7 +1363,7 @@ class AutoMask:
|
|
| 1363 |
"""
|
| 1364 |
self.model = YSAM()
|
| 1365 |
self.model.load_state_dict(
|
| 1366 |
-
state_dict=torch.load(ckpt_path, map_location="cpu")["state_dict"]
|
| 1367 |
)
|
| 1368 |
self.model.eval()
|
| 1369 |
self.model_parallel = torch.nn.DataParallel(self.model)
|
|
|
|
| 1363 |
"""
|
| 1364 |
self.model = YSAM()
|
| 1365 |
self.model.load_state_dict(
|
| 1366 |
+
state_dict=torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
|
| 1367 |
)
|
| 1368 |
self.model.eval()
|
| 1369 |
self.model_parallel = torch.nn.DataParallel(self.model)
|