[DataPipe] Enable profiler record context in __next__ branch

NivekT · pytorchmergebot · commit e8ed16f3c098 · 2022-06-17T16:52:07.000Z
Pull Request resolved: #79757 Approved by: https://github.com/ejguan
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)
 
