[FIX] Error when saving model with custom callbacks (#1487)

Author: marcopeix

neuralforecast/common/_base_model.py

@@ -678,9 +678,20 @@ def on_validation_epoch_end(self):
         self.validation_step_outputs.clear()  # free memory (compute `avg_loss` per epoch)
 
     def save(self, path):
+        import copy
+
+        # Strip callbacks from hparams before saving: callback objects are not
+        # YAML-serializable, which causes PyTorch Lightning to raise a ValueError
+        # during predict() on a loaded model. Callbacks can be re-attached after
+        # loading via `model.trainer_kwargs["callbacks"] = [...]`.
+        # Note: save_hyperparameters() stores **trainer_kwargs contents flat, so
+        # `callbacks` is a top-level key in hparams, not nested under trainer_kwargs.
+        hparams = copy.deepcopy(dict(self.hparams))
+        if "callbacks" in hparams:
+            del hparams["callbacks"]
         with fsspec.open(path, "wb") as f:
             torch.save(
-                {"hyper_parameters": self.hparams, "state_dict": self.state_dict()},
+                {"hyper_parameters": hparams, "state_dict": self.state_dict()},
                 f,
             )
 

tests/test_core.py

@@ -1093,6 +1093,35 @@ def test_save_load_no_dataset(setup_airplane_data):
     np.testing.assert_allclose(forecasts1["DilatedRNN"], forecasts2["DilatedRNN"])
 
 
+def test_save_load_with_callbacks(setup_airplane_data, tmp_path):
+    """Saving a model with trainer callbacks should not break predict after reload.
+
+    Custom callbacks are not YAML-serializable; without the fix this causes a
+    ValueError when PyTorch Lightning's logger tries to log hparams during predict.
+    """
+    from pytorch_lightning.callbacks import Callback
+
+    class _NonYamlCallback(Callback):
+        # lambda attributes are not YAML-safe
+        fn = lambda self: None  # noqa: E731
+
+    AirPassengersPanel_train, _ = setup_airplane_data
+    model = NHITS(
+        h=12,
+        input_size=24,
+        max_steps=10,
+        callbacks=[_NonYamlCallback()],
+    )
+    nf = NeuralForecast(models=[model], freq="M")
+    nf.fit(AirPassengersPanel_train)
+    nf.save(str(tmp_path))
+
+    nf2 = NeuralForecast.load(str(tmp_path))
+    # Should not raise a ValueError from YAML serialization
+    preds = nf2.predict(df=AirPassengersPanel_train)
+    assert preds is not None
+
+
 def test_save_skips_nonzero_ddp_rank(monkeypatch, tmp_path):
     """Only rank 0 should write artifacts when DDP is initialized."""
     dist = pytest.importorskip("torch.distributed")