INNER CODE UNIT · Python

outputs

Adibvafa/CodonTransformer · finetune.py:100

        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 DumpStateDict(pl.callbacks.ModelCheckpoint):
    def __init__(self, checkpoint_dir, checkpoint_filename, every_n_train_steps):
        super().__init__(
            dirpath=checkpoint_dir, every_n_train_steps=every_n_train_steps
        )
        self.checkpoint_filename = checkpoint_filename

View source record →