Only create the data loader iterator when it is being used.
Summary:
Pull Request resolved:
Only create the data loader iterator when it is being used
@@ -252,7 +252,8 @@ class SimpleTrainer(TrainerBase):
self.model = model
self.data_loader = data_loader
- self._data_loader_iter = iter(data_loader)
+ # to access the data loader iterator, call `self._data_loader_iter`
+ self._data_loader_iter_obj = None
self.optimizer = optimizer
def run_step(self):
@@ -293,6 +294,13 @@ class SimpleTrainer(TrainerBase):
"""
self.optimizer.step()
+ @property
+ def _data_loader_iter(self):
+ # only create the data loader iterator when it is used
+ if self._data_loader_iter_obj is None:
+ self._data_loader_iter_obj = iter(self.data_loader)
+ return self._data_loader_iter_obj
+
def _write_metrics(
self,
loss_dict: Mapping[str, torch.Tensor],
