Przeglądaj źródła

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/
Sergii Dymchenko 1 rok temu
rodzic
commit
658aaba327
1 zmienionych plików z 1 dodań i 1 usunięć
  1. 1 1
      sam2/build_sam.py

+ 1 - 1
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)