Fix type annotation for deepspeed training arg (#24988)

a6484c89 · Sylvain Gugger · GitHub · 5b7ffd54 · a6484c89
Unverified Commit a6484c89 authored Jul 21, 2023 by Sylvain Gugger Committed by GitHub Jul 21, 2023
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

src/transformers/training_args.py src/transformers/training_args.py +3 -1

No files found.
--- 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": (