Skip to content

Commit

Permalink
Use weights_only for loading
Browse files Browse the repository at this point in the history
sam2/build_sam.py:81:14: TOR102 [*] `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.

Found with https://github.com/pytorch-labs/torchfix/
  • Loading branch information
kit1980 authored Jul 29, 2024
1 parent 0c5f8c5 commit 658aaba
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion sam2/build_sam.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def build_sam2_video_predictor(

def _load_checkpoint(model, ckpt_path):
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)
if missing_keys:
logging.error(missing_keys)
Expand Down

0 comments on commit 658aaba

Please sign in to comment.