quoccuongle/Grounded-SAM-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #2 from kit1980/patch-1

Browse Source 
Use `weights_only` for loading
This commit is contained in:
Haitham Khedr
2024-10-01 22:32:58 +02:00
committed by GitHub
parent 98fcb164bf 0f6515ae85
commit 52198ead0e
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sam2/build_sam.py
View File

@@ -135,7 +135,7 @@ def build_sam2_video_predictor_hf(model_id, **kwargs):
def _load_checkpoint(model, ckpt_path): def _load_checkpoint(model, ckpt_path):
if ckpt_path is not None: if ckpt_path is not None:
sd = torch.load(ckpt_path, map_location="cpu")["model"] sd = torch.load(ckpt_path, map_location="cpu", weights_only=True)["model"]
missing_keys, unexpected_keys = model.load_state_dict(sd) missing_keys, unexpected_keys = model.load_state_dict(sd)
if missing_keys: if missing_keys:
logging.error(missing_keys) logging.error(missing_keys)