diff --git a/tests/unit/engine/test_engine.py b/tests/unit/engine/test_engine.py index 2ec7186be15..bb1aa36cbba 100644 --- a/tests/unit/engine/test_engine.py +++ b/tests/unit/engine/test_engine.py @@ -194,7 +194,11 @@ def test_exporting(self, fxt_engine, mocker) -> None: checkpoint = "path/to/checkpoint.ckpt" fxt_engine.checkpoint = checkpoint fxt_engine.export() - mock_load_from_checkpoint.assert_called_once_with(checkpoint_path=checkpoint, map_location="cpu") + mock_load_from_checkpoint.assert_called_once_with( + checkpoint_path=checkpoint, + map_location="cpu", + **fxt_engine.model.hparams, + ) mock_export.assert_called_once_with( output_dir=Path(fxt_engine.work_dir), base_name="exported_model",