Fix model saving

huggingface · Jul 12, 2023 · 7d0e1c2 · 7d0e1c2
1 parent 4402d91
commit 7d0e1c2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/optimum/neuron/distributed/base.py b/optimum/neuron/distributed/base.py
@@ -344,6 +344,7 @@ def save_model_checkpoint_as_sharded(
         if get_data_parallel_rank() == 0 and get_tensor_model_parallel_rank() == 0:
             if output_path.is_dir():
                 shutil.rmtree(output_path, ignore_errors=True)
+            output_path.mkdir()
         xm.rendezvous("waiting before saving")
         parallel_layers.save(state_dict, output_path.as_posix())