weights_only=False for omegaconf hparams test

Author: matsumotosan

tests/tests_pytorch/models/test_hparams.py

@@ -19,6 +19,7 @@
 from argparse import Namespace
 from dataclasses import dataclass, field
 from enum import Enum
+from typing import Optional
 from unittest import mock
 
 import cloudpickle
@@ -94,7 +95,9 @@ def __init__(self, hparams, *my_args, **my_kwargs):
 # -------------------------
 # STANDARD TESTS
 # -------------------------
-def _run_standard_hparams_test(tmp_path, model, cls, datamodule=None, try_overwrite=False):
+def _run_standard_hparams_test(
+    tmp_path, model, cls, datamodule=None, try_overwrite=False, weights_only: Optional[bool] = None
+):
     """Tests for the existence of an arg 'test_arg=14'."""
     obj = datamodule if issubclass(cls, LightningDataModule) else model
 
@@ -108,20 +111,20 @@ def _run_standard_hparams_test(tmp_path, model, cls, datamodule=None, try_overwr
 
     # make sure the raw checkpoint saved the properties
     raw_checkpoint_path = _raw_checkpoint_path(trainer)
-    raw_checkpoint = torch.load(raw_checkpoint_path)
+    raw_checkpoint = torch.load(raw_checkpoint_path, weights_only=weights_only)
 
     assert cls.CHECKPOINT_HYPER_PARAMS_KEY in raw_checkpoint
     assert raw_checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY]["test_arg"] == 14
 
     # verify that model loads correctly
-    obj2 = cls.load_from_checkpoint(raw_checkpoint_path)
+    obj2 = cls.load_from_checkpoint(raw_checkpoint_path, weights_only=weights_only)
     assert obj2.hparams.test_arg == 14
 
     assert isinstance(obj2.hparams, hparam_type)
 
     if try_overwrite:
         # verify that we can overwrite the property
-        obj3 = cls.load_from_checkpoint(raw_checkpoint_path, test_arg=78)
+        obj3 = cls.load_from_checkpoint(raw_checkpoint_path, test_arg=78, weights_only=weights_only)
         assert obj3.hparams.test_arg == 78
 
     return raw_checkpoint_path
@@ -176,7 +179,8 @@ def test_omega_conf_hparams(tmp_path, cls):
     assert isinstance(obj.hparams, Container)
 
     # run standard test suite
-    raw_checkpoint_path = _run_standard_hparams_test(tmp_path, model, cls, datamodule=datamodule)
+    # weights_only=False as omegaconf.DictConfig is not an allowed global by default
+    raw_checkpoint_path = _run_standard_hparams_test(tmp_path, model, cls, datamodule=datamodule, weights_only=False)
     obj2 = cls.load_from_checkpoint(raw_checkpoint_path, weights_only=False)
 
     assert isinstance(obj2.hparams, Container)