(self, trainer, pl_module)
| 622 | class CUDACallback(Callback): |
| 623 | # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py |
| 624 | def on_train_epoch_start(self, trainer, pl_module): |
| 625 | # Reset the memory use counter |
| 626 | torch.cuda.reset_peak_memory_stats(trainer.root_gpu) |
| 627 | torch.cuda.synchronize(trainer.root_gpu) |
| 628 | self.start_time = time.time() |
| 629 | |
| 630 | def on_train_epoch_end(self, trainer, pl_module, outputs): |
| 631 | torch.cuda.synchronize(trainer.root_gpu) |
nothing calls this directly
no outgoing calls
no test coverage detected