Unverified Commit 2dc2617c authored by mrshenli's avatar mrshenli Committed by GitHub
Browse files

Adding missing imports in docstring (#760)

`wrap` from `auto_wrap` is used in the docstring example which is missing from the imports.
parent 57821dd2
......@@ -111,7 +111,7 @@ class FullyShardedDataParallel(nn.Module):
across the forward pass. For example::
import torch
from fairscale.nn.auto_wrap import enable_wrap, auto_wrap
from fairscale.nn.auto_wrap import enable_wrap, auto_wrap, wrap
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
with enable_wrap(**fsdp_params):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment