Set weights_only in torch.load (#36991)

2025-03-27 22:55:50 +08:00
parent de77f5b1ec
commit 41a0e58e5b
28 changed files with 64 additions and 78 deletions
--- a/tests/models/videomae/test_modeling_videomae.py
+++ b/tests/models/videomae/test_modeling_videomae.py
@@ -456,7 +456,7 @@ class VideoMAEModelIntegrationTest(unittest.TestCase):

        # add boolean mask, indicating which patches to mask
        local_path = hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos", filename="bool_masked_pos.pt")
-        inputs["bool_masked_pos"] = torch.load(local_path)
+        inputs["bool_masked_pos"] = torch.load(local_path, weights_only=True)

        # forward pass
        with torch.no_grad():