From 658aaba327012309786955f92aa3c15b53fddd39 Mon Sep 17 00:00:00 2001 From: Sergii Dymchenko Date: Mon, 29 Jul 2024 16:54:54 -0700 Subject: [PATCH] Use `weights_only` for loading 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/ --- sam2/build_sam.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sam2/build_sam.py b/sam2/build_sam.py index 39defc46e..06afa5ca4 100644 --- a/sam2/build_sam.py +++ b/sam2/build_sam.py @@ -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)