Fix train-eval mode on trunk

Summary:
train eval broken due to D71827944. The train pipeline automatically raises StopIteration when switching between train mode and eval mode.

This diff sets _data_iter_stopped=False when switching between train mode and eval mode.

Reviewed By: dragonxlwang

Differential Revision: D72582264

fbshipit-source-id: ab9454858a7aa0ec30f96cdab17ed466292951c4

Author: Lucy Liao

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -141,6 +141,11 @@ def __init__(
         self._connected = False
         self._data_iter_stopped = False
 
+    def _reset_data_iter(self) -> None:
+        self._connected = False
+        self._data_iter_stopped = False
+        self._cur_batch = None
+
     def _connect(self, dataloader_iter: Iterator[In]) -> None:
         cur_batch = next(dataloader_iter)
         self._cur_batch = cur_batch