| 585 | return False |
| 586 | |
| 587 | def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): |
| 588 | if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): |
| 589 | self.log_img(pl_module, batch, batch_idx, split="train") |
| 590 | # if self.save_freq is not None: |
| 591 | # epoch = trainer.current_epoch |
| 592 | # global_step = trainer.global_step |
| 593 | # if global_step % self.save_freq == 0: |
| 594 | # filename = f'{epoch}_{global_step}.ckpt' |
| 595 | # ckpt_path = os.path.join(trainer.checkpoint_callback.dirpath, filename) |
| 596 | # trainer.save_checkpoint(ckpt_path) |
| 597 | |
| 598 | |
| 599 | class CUDACallback(Callback): |