Fixed failing tests on pytorch nightly using torch.load

pytorch · Nov 7, 2024 · c50b777 · c50b777
1 parent 3c5e213
commit c50b777
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/tests/ignite/engine/test_deterministic.py b/tests/ignite/engine/test_deterministic.py
@@ -10,6 +10,7 @@
 import torch.nn as nn
 from torch.optim import SGD
 from torch.utils.data import BatchSampler, DataLoader, RandomSampler
+from packaging.version import Version
 
 import ignite.distributed as idist
 from ignite.engine import Events
@@ -737,7 +738,11 @@ def write_data_grads_weights(e):
             grad_norms.append([i, total[1]] + out2)
 
         if sd is not None:
-            sd = torch.load(sd)
+            if Version(torch.__version__) >= Version("1.13.0"):
+                kwargs = {"weights_only": False}
+            else:
+                kwargs = {}
+            sd = torch.load(sd, **kwargs)
             model.load_state_dict(sd[0])
             opt.load_state_dict(sd[1])
             from ignite.engine.deterministic import _repr_rng_state

diff --git a/tests/ignite/handlers/test_state_param_scheduler.py b/tests/ignite/handlers/test_state_param_scheduler.py
@@ -295,7 +295,11 @@ def test_torch_save_load(dirname):
 
     filepath = Path(dirname) / "dummy_lambda_state_parameter_scheduler.pt"
     torch.save(lambda_state_parameter_scheduler, filepath)
-    loaded_lambda_state_parameter_scheduler = torch.load(filepath)
+    if Version(torch.__version__) >= Version("1.13.0"):
+        kwargs = {"weights_only": False}
+    else:
+        kwargs = {}
+    loaded_lambda_state_parameter_scheduler = torch.load(filepath, **kwargs)
 
     engine1 = Engine(lambda e, b: None)
     lambda_state_parameter_scheduler.attach(engine1, Events.EPOCH_COMPLETED)