Fix unit test

openvinotoolkit · Jun 12, 2024 · 157595b · 157595b
1 parent 227e0be
commit 157595b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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",