add testing

Author: SkafteNicki

tests/tests_pytorch/callbacks/test_stochastic_weight_avg.py

@@ -387,5 +387,35 @@ def test_misconfiguration_error_with_sharded_model(tmp_path, strategy: str):
         trainer.fit(model)
 
 
+def test_swa_with_infinite_epochs_and_batchnorm(tmp_path):
+    """Test that SWA works correctly with max_epochs=-1 (infinite training) and BatchNorm."""
+    model = SwaTestModel(batchnorm=True)
+    swa_callback = StochasticWeightAveraging(swa_lrs=0.1, swa_epoch_start=2)
+
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        max_epochs=-1,
+        max_steps=30,  # Use max_steps as stopping condition
+        limit_train_batches=5,
+        limit_val_batches=0,
+        callbacks=[swa_callback],
+        logger=False,
+    )
+    assert trainer.max_epochs == -1
+    assert trainer.fit_loop.max_epochs == -1
+
+    trainer.fit(model)
+    assert trainer.current_epoch >= 5
+    assert trainer.global_step == 30
+    assert trainer.max_epochs == -1
+
+    # Verify SWA was actually applied (update_parameters should have been called)
+    # SWA starts at epoch 2, so with 6 epochs (0-5), we should have 4 updates (epochs 2, 3, 4, 5)
+    assert swa_callback.n_averaged is not None
+    assert swa_callback.n_averaged > 0, "SWA should have updated parameters"
+
+
 def _backward_patch(trainer: Trainer) -> AbstractContextManager:
     return mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward)