From 0b087eb81592a2a2948680cce1b0afaf69b345d8 Mon Sep 17 00:00:00 2001 From: v-chen_data Date: Thu, 6 Jun 2024 14:42:19 -0700 Subject: [PATCH] quick documentation fix --- composer/distributed/dist_strategy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/composer/distributed/dist_strategy.py b/composer/distributed/dist_strategy.py index ad08e172b8..1cc1044a02 100644 --- 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.