frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[DataPipe] Enable profiler record context in __next__ branch

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

Approved by: https://github.com/ejguan
This commit is contained in:
Kevin Tse
2022-06-16 18:53:15 -04:00
committed by PyTorch MergeBot
parent 25ca006707
commit e8ed16f3c0
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torch/utils/data/datapipes/_hook_iterator.py
View File

@ -171,7 +171,8 @@ def hook_iterator(namespace, profile_name):
@functools.wraps(next_func) @functools.wraps(next_func)
def wrap_next(*args, **kwargs): def wrap_next(*args, **kwargs):
if torch.autograd._profiler_enabled(): if torch.autograd._profiler_enabled():
return next_func(*args, **kwargs) with profiler_record_fn_context():
return next_func(*args, **kwargs)
else: else:
return next_func(*args, **kwargs) return next_func(*args, **kwargs)