murcherful commited on
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)