From 99c846b175d41207a8295525424a4991be211873 Mon Sep 17 00:00:00 2001 From: Joshua David Date: Tue, 9 Jul 2024 22:28:44 -0700 Subject: [PATCH] Save the final model --- train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/train.py b/train.py index de5756e..22c324a 100644 --- a/train.py +++ b/train.py @@ -470,6 +470,10 @@ def main(): logger.info(f"Validation loss after short context recovery: {val_loss:.4f}") wandb.log({"short_context_val_loss": val_loss}) + # Save the final model + accelerator.save_state("final_model.pt") + wandb.save("final_model.pt") + # Finish logging and close the Weights & Biases run wandb.finish()