Method on_train_batch_end

train_pl.py:115–116

(self, *args, **kwargs)

        return loss

    def on_train_batch_end(self, *args, **kwargs):
        update_ema(self.ema, self.model)

    def on_save_checkpoint(self, checkpoint):
        super().on_save_checkpoint(checkpoint)