onnxruntime
4e377a13 - Use weights_only=True for remaining torch.load() calls (#28421)

Commit
8 days ago
Use weights_only=True for remaining torch.load() calls (#28421) ### Description Follow-up to PR #28097. Applies the same `_torch_load_weights_only()` wrapper to the two remaining `torch.load()` call sites. `torch.load` can deserialize arbitrary Python pickle payloads. Using `weights_only=True` restricts loading to tensor/checkpoint data on supported PyTorch versions and is the safer default. The wrapper gracefully falls back to the default `torch.load` behavior on older PyTorch versions that do not support the `weights_only` parameter. ### Summary of Changes | File | Change | |------|--------| | `onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py` | Adds `_torch_load_weights_only()` helper and uses it when loading scheduler/optimizer state dicts. | | `orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py` | Adds `_torch_load_weights_only()` helper and uses it when loading DDP model checkpoint. | ### Motivation and Context These were the last two `torch.load()` calls in the repository without `weights_only=True`. While both are in test/tooling code with low direct risk, this change ensures consistency with the pattern established in PR #28097 and eliminates all unsafe deserialization call sites. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Parents
Loading