Fix failing test

huggingface · Oct 31, 2024 · 1bc5a50 · 1bc5a50
1 parent 3146290
commit 1bc5a50
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
@@ -264,6 +264,7 @@ def test_save_torch_model(mocker: MockerFixture, tmp_path: Path) -> None:
         max_shard_size="3GB",
         metadata={"foo": "bar"},
         safe_serialization=True,
+        is_main_process=True,
     )
     safe_state_dict_mock.assert_called_once_with(
         state_dict=model_mock.state_dict.return_value,
@@ -273,6 +274,7 @@ def test_save_torch_model(mocker: MockerFixture, tmp_path: Path) -> None:
         max_shard_size="3GB",
         metadata={"foo": "bar"},
         safe_serialization=True,
+        is_main_process=True,
     )