INNER CODE UNIT · Python
on_train_epoch_end
Adibvafa/CodonTransformer · pretrain.py:118
def on_train_epoch_end(self, trainer, pl_module):
current_epoch = trainer.current_epoch
if current_epoch % self.save_interval == 0 or current_epoch == 0:
checkpoint_path = os.path.join(
self.checkpoint_dir, f"epoch_{current_epoch}.ckpt"
)
trainer.save_checkpoint(checkpoint_path)
print(f"\nCheckpoint saved at {checkpoint_path}\n")
def main(args):
"""Pretrain the CodonTransformer model."""
pl.seed_everything(args.seed)
torch.set_float32_matmul_precision("medium")
# Load the tokenizer and model
tokenizer = PreTrainedTokenizerFast(
tokenizer_file=args.tokenizer_path,