From 3fe24427ab8fb13aa9d033ae533a9c82ee3daa40 Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Tue, 17 Oct 2023 10:29:58 -0700 Subject: [PATCH] Fix prepare_module docstring Summary: Created from CodeHub with https://fburl.com/edit-in-codehub Reviewed By: ananthsub Differential Revision: D50369401 --- torchtnt/utils/prepare_module.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/torchtnt/utils/prepare_module.py b/torchtnt/utils/prepare_module.py index 15eda97532..86eb156de2 100644 --- a/torchtnt/utils/prepare_module.py +++ b/torchtnt/utils/prepare_module.py @@ -275,7 +275,6 @@ def prepare_module( module: module to be used. device: device to which module will be moved. strategy: the data parallelization strategy to be used. if a string, must be one of ``ddp``, ``fsdp``, or ``noop``. - swa_params: params for stochastic weight averaging https://pytorch.org/docs/stable/optim.html#stochastic-weight-averaging. torch_compile_params: params for Torch compile https://pytorch.org/docs/stable/generated/torch.compile.html. activation_checkpoint_params: params for enabling activation checkpointing. """ @@ -321,7 +320,7 @@ def prepare_module( module.compile(**asdict(torch_compile_params)) except AttributeError: rank_zero_warn( - "Please install pytorch nightlies to use in-place compile to avoid altering the state_dict keys when checkpointing." + "Please install PyTorch nightlies to use in-place compile to avoid altering the state_dict keys when checkpointing." ) torch.compile(module, **asdict(torch_compile_params))