(
self,
trainer: Trainer,
pl_module: LightningModule,
outputs: Any,
batch: Any,
batch_idx: int,
)
| 387 | |
| 388 | @trainer_rank_zero_only |
| 389 | def on_train_batch_end( |
| 390 | self, |
| 391 | trainer: Trainer, |
| 392 | pl_module: LightningModule, |
| 393 | outputs: Any, |
| 394 | batch: Any, |
| 395 | batch_idx: int, |
| 396 | ) -> None: |
| 397 | self.total_lengths += self.length_fn(batch) |
| 398 | if trainer.fit_loop._should_accumulate(): |
| 399 | return |
| 400 | train_elapsed = time.perf_counter() - self.train_t0 |
| 401 | assert self.speed_monitor is not None |
| 402 | iter_num = trainer.fit_loop.total_batch_idx |
| 403 | assert (measured_flops := pl_module.measured_flops) is not None |
| 404 | self.speed_monitor.on_train_batch_end( |
| 405 | (iter_num + 1) * self.batch_size, |
| 406 | train_elapsed, |
| 407 | # this assumes that device FLOPs are the same and that all devices have the same batch size |
| 408 | trainer.world_size, |
| 409 | flops_per_batch=measured_flops, |
| 410 | lengths=self.total_lengths, |
| 411 | ) |
| 412 | |
| 413 | @trainer_rank_zero_only |
| 414 | def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None: |
nothing calls this directly
no test coverage detected