quick documentation fix

mosaicml · Jun 6, 2024 · 0b087eb · 0b087eb
1 parent 07d5d3a
commit 0b087eb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/composer/distributed/dist_strategy.py b/composer/distributed/dist_strategy.py
@@ -209,7 +209,7 @@ def prepare_fsdp_module(
     Args:
         model (torch.nn.Module): The model to wrap.
         optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional): The optimizer for `model`, assumed to have a single param group := model.parameters().
-        fsdp_config (dict[str, Any]): The FSDP config.
+        fsdp_config (FSDPConfig): The FSDP config.
         precision: (Precision): The precision being used by the Trainer, used to fill in defaults for FSDP `mixed_precision` settings.
         device (Device): The device being used by the Trainer.
         auto_microbatching (bool, optional): Whether or not auto microbatching is enabled.