Unverified Commit 02a8913c authored by tmarkstrum's avatar tmarkstrum Committed by GitHub
Browse files

fix trailing space issue (#903)

parent d3417ceb
...@@ -2366,8 +2366,8 @@ def _post_state_dict_hook( ...@@ -2366,8 +2366,8 @@ def _post_state_dict_hook(
prefix: str, prefix: str,
*args: Any, *args: Any,
) -> "OrderedDict[str, torch.Tensor]": ) -> "OrderedDict[str, torch.Tensor]":
# When state_dict_on_rank_0_only is ``True``, ``model.state_dict()`` will only # When state_dict_on_rank_0_only is ``True``, ``model.state_dict()`` will only
# returns full state dict on rank 0 and return empty dict non-rank 0, # returns full state dict on rank 0 and return empty dict non-rank 0,
# which allow FullyShardedDataParallel to skip the GPU -> CPU copy on # which allow FullyShardedDataParallel to skip the GPU -> CPU copy on
# non-rank 0 altogether and prevent OOM. # non-rank 0 altogether and prevent OOM.
if state_dict_on_rank_0_only and dist.get_rank() != 0: if state_dict_on_rank_0_only and dist.get_rank() != 0:
......
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