We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 51f640e commit 5b68901Copy full SHA for 5b68901
intermediate_source/FSDP_tutorial.rst
@@ -217,8 +217,8 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
217
train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
218
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
219
my_auto_wrap_policy = functools.partial(
220
- size_based_auto_wrap_policy, min_num_params=100
221
- )
+ size_based_auto_wrap_policy, min_num_params=100
+ )
222
torch.cuda.set_device(rank)
223
224
0 commit comments