weights_only=False for torch>=2.6 in tests

Author: matsumotosan

tests/tests_pytorch/callbacks/test_callbacks.py

@@ -18,6 +18,7 @@
 import pytest
 from lightning_utilities.test.warning import no_warning_call
 
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.pytorch import Callback, Trainer
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.demos.boring_classes import BoringModel
@@ -132,7 +133,8 @@ def test_resume_callback_state_saved_by_type_stateful(tmp_path):
 
     callback = OldStatefulCallback(state=222)
     trainer = Trainer(default_root_dir=tmp_path, max_steps=2, callbacks=[callback])
-    trainer.fit(model, ckpt_path=ckpt_path)
+    weights_only = False if _TORCH_GREATER_EQUAL_2_6 else None
+    trainer.fit(model, ckpt_path=ckpt_path, weights_only=weights_only)
     assert callback.state == 111
 
 

tests/tests_pytorch/callbacks/test_stochastic_weight_avg.py

@@ -25,6 +25,7 @@
 from torch.optim.swa_utils import SWALR
 from torch.utils.data import DataLoader
 
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import StochasticWeightAveraging
 from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset, RandomIterableDataset
@@ -173,8 +174,9 @@ def train_with_swa(
         devices=devices,
     )
 
+    weights_only = False if _TORCH_GREATER_EQUAL_2_6 else None
     with _backward_patch(trainer):
-        trainer.fit(model)
+        trainer.fit(model, weights_only=weights_only)
 
     # check the model is the expected
     assert trainer.lightning_module == model
@@ -307,8 +309,9 @@ def _swa_resume_training_from_checkpoint(tmp_path, model, resume_model, ddp=Fals
     }
     trainer = Trainer(callbacks=SwaTestCallback(swa_epoch_start=swa_start, swa_lrs=0.1), **trainer_kwargs)
 
+    weights_only = False if _TORCH_GREATER_EQUAL_2_6 else None
     with _backward_patch(trainer), pytest.raises(Exception, match="SWA crash test"):
-        trainer.fit(model)
+        trainer.fit(model, weights_only=weights_only)
 
     checkpoint_dir = Path(tmp_path) / "checkpoints"
     checkpoint_files = os.listdir(checkpoint_dir)
@@ -318,7 +321,7 @@ def _swa_resume_training_from_checkpoint(tmp_path, model, resume_model, ddp=Fals
     trainer = Trainer(callbacks=SwaTestCallback(swa_epoch_start=swa_start, swa_lrs=0.1), **trainer_kwargs)
 
     with _backward_patch(trainer):
-        trainer.fit(resume_model, ckpt_path=ckpt_path)
+        trainer.fit(resume_model, ckpt_path=ckpt_path, weights_only=weights_only)
 
 
 class CustomSchedulerModel(SwaTestModel):

tests/tests_pytorch/models/test_hparams.py

@@ -30,6 +30,7 @@
 from lightning_utilities.test.warning import no_warning_call
 from torch.utils.data import DataLoader
 
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.core.datamodule import LightningDataModule
@@ -748,8 +749,9 @@ def test_model_with_fsspec_as_parameter(tmp_path):
     trainer = Trainer(
         default_root_dir=tmp_path, limit_train_batches=2, limit_val_batches=2, limit_test_batches=2, max_epochs=1
     )
-    trainer.fit(model)
-    trainer.test()
+    weights_only = False if _TORCH_GREATER_EQUAL_2_6 else None
+    trainer.fit(model, weights_only=weights_only)
+    trainer.test(weights_only=weights_only)
 
 
 @pytest.mark.xfail(