def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
    self.initial_target_steps = lightning_getattr(pl_module, self.target_steps_attr)
    lightning_setattr(pl_module, self.target_steps_attr, self.initial_steps)
    trainer.reset_train_dataloader(pl_module)
def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
    self.initial_target_steps = lightning_getattr(pl_module, self.target_steps_attr)
    lightning_setattr(pl_module, self.target_steps_attr, self.initial_steps)
    trainer.reset_train_dataloader(pl_module)