| from pytorch_lightning.callbacks import Callback | |
| class IncreaseDataEpoch(Callback): | |
| def __init__(self): | |
| super().__init__() | |
| def on_train_epoch_start(self, trainer, pl_module): | |
| epoch = pl_module.current_epoch | |
| if hasattr(trainer.datamodule.train_dataset, "shared_epoch"): | |
| trainer.datamodule.train_dataset.shared_epoch.set_value(epoch) | |