Revert "try drop pickle warning"

This reverts commit 80a089d.

Author: Borda

tests/tests_pytorch/callbacks/test_early_stopping.py

@@ -193,11 +193,13 @@ def test_pickling():
     early_stopping = EarlyStopping(monitor="foo")
 
     early_stopping_pickled = pickle.dumps(early_stopping)
-    early_stopping_loaded = pickle.loads(early_stopping_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        early_stopping_loaded = pickle.loads(early_stopping_pickled)
     assert vars(early_stopping) == vars(early_stopping_loaded)
 
     early_stopping_pickled = cloudpickle.dumps(early_stopping)
-    early_stopping_loaded = cloudpickle.loads(early_stopping_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        early_stopping_loaded = cloudpickle.loads(early_stopping_pickled)
     assert vars(early_stopping) == vars(early_stopping_loaded)
 
 

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -352,11 +352,13 @@ def test_pickling(tmp_path):
     ckpt = ModelCheckpoint(dirpath=tmp_path)
 
     ckpt_pickled = pickle.dumps(ckpt)
-    ckpt_loaded = pickle.loads(ckpt_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        ckpt_loaded = pickle.loads(ckpt_pickled)
     assert vars(ckpt) == vars(ckpt_loaded)
 
     ckpt_pickled = cloudpickle.dumps(ckpt)
-    ckpt_loaded = cloudpickle.loads(ckpt_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        ckpt_loaded = cloudpickle.loads(ckpt_pickled)
     assert vars(ckpt) == vars(ckpt_loaded)
 
 

tests/tests_pytorch/core/test_metric_result_integration.py

@@ -254,7 +254,8 @@ def lightning_log(fx, *args, **kwargs):
         }
 
         # make sure can be pickled
-        pickle.loads(pickle.dumps(result))
+        with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+            pickle.loads(pickle.dumps(result))
         # make sure can be torch.loaded
         filepath = str(tmp_path / "result")
         torch.save(result, filepath)

tests/tests_pytorch/helpers/test_datasets.py

@@ -44,7 +44,9 @@ def test_pickling_dataset_mnist(dataset_cls, args):
     mnist = dataset_cls(**args)
 
     mnist_pickled = pickle.dumps(mnist)
-    pickle.loads(mnist_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        pickle.loads(mnist_pickled)
 
     mnist_pickled = cloudpickle.dumps(mnist)
-    cloudpickle.loads(mnist_pickled)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        cloudpickle.loads(mnist_pickled)

tests/tests_pytorch/loggers/test_all.py

@@ -184,7 +184,12 @@ def _test_loggers_pickle(tmp_path, monkeypatch, logger_class: Logger):
     trainer = Trainer(max_epochs=1, logger=logger)
     pkl_bytes = pickle.dumps(trainer)
 
-    trainer2 = pickle.loads(pkl_bytes)
+    with (
+        pytest.warns(FutureWarning, match="`weights_only=False`")
+        if _TORCH_EQUAL_2_4_0 or (_TORCH_GREATER_EQUAL_2_4_1 and logger_class not in (CSVLogger, TensorBoardLogger))
+        else nullcontext()
+    ):
+        trainer2 = pickle.loads(pkl_bytes)
     trainer2.logger.log_metrics({"acc": 1.0})
 
     # make sure we restored properly

tests/tests_pytorch/loggers/test_logger.py

@@ -124,7 +124,8 @@ def test_multiple_loggers_pickle(tmp_path):
 
     trainer = Trainer(logger=[logger1, logger2])
     pkl_bytes = pickle.dumps(trainer)
-    trainer2 = pickle.loads(pkl_bytes)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        trainer2 = pickle.loads(pkl_bytes)
     for logger in trainer2.loggers:
         logger.log_metrics({"acc": 1.0}, 0)
 

tests/tests_pytorch/loggers/test_wandb.py

@@ -162,7 +162,8 @@ def name(self):
     assert trainer.logger.experiment, "missing experiment"
     assert trainer.log_dir == logger.save_dir
     pkl_bytes = pickle.dumps(trainer)
-    trainer2 = pickle.loads(pkl_bytes)
+    with pytest.warns(FutureWarning, match="`weights_only=False`") if _TORCH_EQUAL_2_4_0 else nullcontext():
+        trainer2 = pickle.loads(pkl_bytes)
 
     assert os.environ["WANDB_MODE"] == "dryrun"
     assert trainer2.logger.__class__.__name__ == WandbLogger.__name__