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/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -648,7 +648,7 @@ class TrainerIntegrationCommon:
        else:
            best_model = RegressionModel()
            if not safe_weights:
-                state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME))
+                state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME), weights_only=True)
            else:
                state_dict = safetensors.torch.load_file(os.path.join(checkpoint, SAFE_WEIGHTS_NAME))
            best_model.load_state_dict(state_dict)