Save the final model

jshuadvd · Jul 10, 2024 · 99c846b · 99c846b
1 parent 7f2c9ef
commit 99c846b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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()