fix mock (#872)

mosaicml · Jan 14, 2024 · b40b805 · b40b805
1 parent b69318e
commit b40b805
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -300,6 +300,8 @@ def test_huggingface_conversion_callback_interval(
     mlflow_logger_mock.save_model = MagicMock()
     mlflow_logger_mock.register_model = MagicMock()
     mlflow_logger_mock.model_registry_prefix = ''
+    mlflow_logger_mock._experiment_id = 'mlflow-experiment-id'
+    mlflow_logger_mock._run_id = 'mlflow-run-id'
     trainer = Trainer(
         model=original_model,
         device='gpu',
@@ -534,6 +536,8 @@ def test_huggingface_conversion_callback(
     mlflow_logger_mock.save_model = MagicMock()
     mlflow_logger_mock.register_model = MagicMock()
     mlflow_logger_mock.model_registry_prefix = ''
+    mlflow_logger_mock._experiment_id = 'mlflow-experiment-id'
+    mlflow_logger_mock._run_id = 'mlflow-run-id'
     trainer = Trainer(
         model=original_model,
         device='gpu',