add label for data loader (#1907)

Summary:

The data loader batch fetching logic isn't regular kernel related so profiler wouldn't show the actual e2e time.

This would show a label instead of just leaving it blank.

Differential Revision: D55858614

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -283,7 +283,8 @@ def _next_batch(self, dataloader_iter: Iterator[In]) -> Optional[In]:
         if self._dataloader_exhausted:
             batch = None
         else:
-            batch = next(dataloader_iter, None)
+            with record_function("## next_batch ##"):
+                batch = next(dataloader_iter, None)
             if batch is None:
                 self._dataloader_exhausted = True
         return batch
@@ -700,7 +701,8 @@ def _next_batch(self, dataloader_iter: Iterator[In]) -> Optional[In]:
         if self._dataloader_exhausted:
             batch = None
         else:
-            batch = next(dataloader_iter, None)
+            with record_function("## next_batch ##"):
+                batch = next(dataloader_iter, None)
             if batch is None:
                 self._dataloader_exhausted = True
         return batch