[DataPipe] Enable profiler record context in __next__ branch (#79757)

NivekT · facebook-github-bot · commit 453a16918dbf · 2022-06-20T08:07:31.000-07:00
Summary: Pull Request resolved: #79757 Approved by: https://github.com/ejguan Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/e8ed16f3c098290b52accfcc07ea668678ac2dab Reviewed By: malfet Differential Revision: D37258847 Pulled By: NivekT fbshipit-source-id: e011c12e522766f3f1234d5d00cfdd22aa72954b
diff --git a/torch/utils/data/datapipes/_hook_iterator.py b/torch/utils/data/datapipes/_hook_iterator.py
@@ -171,7 +171,8 @@ def wrap_generator(*args, **kwargs):
             @functools.wraps(next_func)
             def wrap_next(*args, **kwargs):
                 if torch.autograd._profiler_enabled():
-                    return next_func(*args, **kwargs)
+                    with profiler_record_fn_context():
+                        return next_func(*args, **kwargs)
                 else:
                     return next_func(*args, **kwargs)
 
