more typing

mosaicml · Sep 20, 2024 · 2390ae6 · 2390ae6
1 parent c1d6397
commit 2390ae6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/composer/utils/parallelism.py b/composer/utils/parallelism.py
@@ -30,7 +30,7 @@ class FSDPConfig:
     limit_all_gathers: bool = True
     load_monolith_rank0_only: bool = False
     load_planner: Optional[Any] = None
-    mixed_precision: Union[str, dict[str, str]] = 'DEFAULT'
+    mixed_precision: str = 'DEFAULT'
     process_group: Optional[Any] = None
     save_planner: Optional[Any] = None
     sharded_ckpt_prefix_dir: str = 'ep{epoch}-ba{batch}'