pytorch
e7efeb52 - [FSDP] Save `_stream_to_name` for debugging (#90611)

Commit
3 years ago
[FSDP] Save `_stream_to_name` for debugging (#90611) This saves a data structure `_stream_to_name: Dict[torch.cuda.Stream, str]` that maps each FSDP stream to its name. This can help in debugging by checking `_stream_to_name[torch.cuda.current_stream()]` to see if it is `"default"` or `"unshard"` in the post-backward hook for example. Pull Request resolved: https://github.com/pytorch/pytorch/pull/90611 Approved by: https://github.com/rohan-varma
Author
Andrew Gu
Committer
Parents
Loading