INNER CODE UNIT · Python

training_step

Adibvafa/CodonTransformer · pretrain.py:98

    def training_step(self, batch, batch_idx):
        self.model.bert.set_attention_type("block_sparse")
        outputs = self.model(**batch)
        self.log_dict(
            dictionary={
                "loss": outputs.loss,
                "lr": self.trainer.optimizers[0].param_groups[0]["lr"],
            },
            on_step=True,
            prog_bar=True,
        )
        return outputs.loss


class EpochCheckpoint(pl.Callback):
    def __init__(self, checkpoint_dir, save_interval):
        super().__init__()
        self.checkpoint_dir = checkpoint_dir

View source record →