OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[FSDP][Easy] Fix state_dict_type() docstring example

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77848

Approved by: https://github.com/rohan-varma
This commit is contained in:
Andrew Gu 2022-05-19 20:36:46 +00:00 committed by PyTorch MergeBot
parent 96e674a0c9
commit 36bf8007f7
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
torch/distributed/fsdp/fully_sharded_data_parallel.py
View File

@ -1708,8 +1708,7 @@ class FullyShardedDataParallel(nn.Module):
Example::
>>> model = DDP(FSDP(...))
>>> fsdp_root = model.module
>>> with FSDP.state_dict_type(fsdp_root, StateDictType.LOCAL_STATE_DICT):
>>> with FSDP.state_dict_type(model, StateDictType.LOCAL_STATE_DICT):
>>> checkpoint = model.state_dict()
Args: