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/utils/test_modeling_utils.py
+++ b/tests/utils/test_modeling_utils.py
@@ -765,7 +765,7 @@ class ModelUtilsTest(TestCasePlus):
                    # Note: pickle adds some junk so the weight of the file can end up being slightly bigger than
                    # the size asked for (since we count parameters)
                    if size >= max_size_int + 50000:
-                        state_dict = torch.load(shard_file)
+                        state_dict = torch.load(shard_file, weights_only=True)
                        self.assertEqual(len(state_dict), 1)

                # Check the index and the shard files found match