Fix type annotation for deepspeed training arg (huggingface#24988)

blbadger · Nov 8, 2023 · b93c8bb · b93c8bb
1 parent 0766e66
commit b93c8bb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
@@ -1001,6 +1001,7 @@ class TrainingArguments:
             )
         },
     )
+    # Do not touch this type annotation or it will stop working in CLI
     fsdp_config: Optional[str] = field(
         default=None,
         metadata={
@@ -1019,7 +1020,8 @@ class TrainingArguments:
             )
         },
     )
-    deepspeed: Optional[Union[str, Dict]] = field(
+    # Do not touch this type annotation or it will stop working in CLI
+    deepspeed: Optional[str] = field(
         default=None,
         metadata={
             "help": (