| import time | |
| import pytorch_lightning as pl | |
| class TimeLimitCallback(pl.Callback): | |
| def __init__(self, max_duration_seconds): | |
| super().__init__() | |
| self.max_duration_seconds = max_duration_seconds | |
| self.start_time = None | |
| def on_train_start(self, trainer, pl_module): | |
| self.start_time = time.time() | |
| def on_batch_end(self, trainer, pl_module): | |
| elapsed_time = time.time() - self.start_time | |
| if elapsed_time >= self.max_duration_seconds: | |
| print(f"Training stopped due to time constraint ({self.max_duration_seconds} seconds).") | |
| trainer.should_stop = True # Stops training | |