try patching cpu count for test

mosaicml · Oct 31, 2023 · d54da2a · d54da2a
1 parent b4fc316
commit d54da2a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
@@ -5,7 +5,7 @@
 import os
 import pathlib
 import sys
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from composer import Trainer
 from composer.loggers import MLFlowLogger
@@ -254,6 +254,7 @@ def test_callback_inits_with_defaults():
 @pytest.mark.parametrize(
     'hf_save_interval,save_interval,max_duration,expected_hf_checkpoints,expected_normal_checkpoints',
     [('3ba', '2ba', '7ba', 3, 4), ('1dur', '2ba', '1ep', 1, 4)])
+@patch('os.cpu_count', MagicMock(return_value=None))
 def test_huggingface_conversion_callback(model: str, tmp_path: pathlib.Path,
                                          fsdp_state_dict_type: Optional[str],
                                          log_to_mlflow: bool,