Add a simple validation step after short context recovery

jshuadvd · Jul 10, 2024 · 7f2c9ef · 7f2c9ef
1 parent 118c02e
commit 7f2c9ef
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -460,6 +460,16 @@ def main():
                 tokenizer=tokenizer,
             )
 
+        # Add a simple validation step after short context recovery
+        model.eval()
+        with torch.no_grad():
+            val_loss = sum(
+                criterion(model(inputs), targets).item()
+                for inputs, targets in val_loader
+            ) / len(val_loader)
+        logger.info(f"Validation loss after short context recovery: {val_loss:.4f}")
+        wandb.log({"short_context_val_loss": val_loss})
+
     # Finish logging and close the Weights & Biases run
     wandb.finish()