Update the train method to use the passkey test

jshuadvd · Jul 5, 2024 · 1400108 · 1400108
1 parent b601dcf
commit 1400108
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -335,7 +335,15 @@ def main():
     scheduler = CosineAnnealingLR(optimizer, T_max=10)
 
     # Train the recovered model
-    train(recovered_model, train_loader, val_loader, optimizer, criterion, scheduler)
+    train(
+        recovered_model,
+        train_loader,
+        val_loader,
+        optimizer,
+        criterion,
+        scheduler,
+        tokenizer,
+    )
 
     # Finish logging and close the Weights & Biases run
     wandb.finish()