Use weights_only=True for remaining torch.load() calls

onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py

@@ -4,9 +4,28 @@
 """This file is used to generate test data for LR scheduler optimizer tests in
 orttraining/orttraining/test/training_api/core/training_api_tests.cc."""
 
+import inspect
+import logging
+
 import torch
 from torch.optim.lr_scheduler import LambdaLR
 
+logger = logging.getLogger(__name__)
+
+_TORCH_LOAD_HAS_WEIGHTS_ONLY = "weights_only" in inspect.signature(torch.load).parameters
+
+
+def _torch_load_weights_only(path: str, **kwargs):
+    if _TORCH_LOAD_HAS_WEIGHTS_ONLY:
+        return torch.load(path, weights_only=True, **kwargs)
+
+    logger.warning(
+        "Current PyTorch version does not support torch.load(..., weights_only=True); "
+        "falling back to default torch.load behavior for %s.",
+        path,
+    )
+    return torch.load(path, **kwargs)
+
 
 class SingleParameterModule(torch.nn.Module):
     """A dummy module containing only one trainable parameter."""
@@ -90,7 +109,7 @@ def main():
             json.dump(data, f, ensure_ascii=False, indent=4)
 
         data = []
-        state_dict = torch.load(fp.name)
+        state_dict = _torch_load_weights_only(fp.name)
         new_adamw_optimizer = torch.optim.AdamW(pt_model.parameters(), lr=1e-3)
         new_adamw_optimizer.load_state_dict(state_dict["optimizer"])
 

orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py

@@ -1,6 +1,8 @@
 # This test script is a modified version of Pytorch's tutorial.
 # For details, see https://pytorch.org/tutorials/intermediate/ddp_tutorial.html.
 import argparse
+import inspect
+import logging
 import os
 import sys  # noqa: F401
 import tempfile
@@ -15,6 +17,22 @@
 import onnxruntime  # noqa: F401
 from onnxruntime.training.ortmodule import ORTModule
 
+logger = logging.getLogger(__name__)
+
+_TORCH_LOAD_HAS_WEIGHTS_ONLY = "weights_only" in inspect.signature(torch.load).parameters
+
+
+def _torch_load_weights_only(path: str, **kwargs):
+    if _TORCH_LOAD_HAS_WEIGHTS_ONLY:
+        return torch.load(path, weights_only=True, **kwargs)
+
+    logger.warning(
+        "Current PyTorch version does not support torch.load(..., weights_only=True); "
+        "falling back to default torch.load behavior for %s.",
+        path,
+    )
+    return torch.load(path, **kwargs)
+
 
 def setup(rank, world_size):
     os.environ["MASTER_ADDR"] = "localhost"
@@ -113,7 +131,7 @@ def demo_checkpoint(rank, world_size, use_ort_module):
     dist.barrier()
     # configure map_location properly
     map_location = {"cuda:0": f"cuda:{rank}"}
-    ddp_model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=map_location))
+    ddp_model.load_state_dict(_torch_load_weights_only(CHECKPOINT_PATH, map_location=map_location))
 
     optimizer.zero_grad()
     outputs = ddp_model(torch.randn(20, 10))