| 608 | return False |
| 609 | |
| 610 | def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): |
| 611 | if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): |
| 612 | self.log_img(pl_module, batch, batch_idx, split="train") |
| 613 | # if self.save_freq is not None: |
| 614 | # epoch = trainer.current_epoch |
| 615 | # global_step = trainer.global_step |
| 616 | # if global_step % self.save_freq == 0: |
| 617 | # filename = f'{epoch}_{global_step}.ckpt' |
| 618 | # ckpt_path = os.path.join(trainer.checkpoint_callback.dirpath, filename) |
| 619 | # trainer.save_checkpoint(ckpt_path) |
| 620 | |
| 621 | |
| 622 | class CUDACallback(Callback): |