From 49efa5cb4007292cf2f0a6436e3fbc3fd71fce4e Mon Sep 17 00:00:00 2001
From: Jake Poznanski <jakep@allenai.org>
Date: Thu, 26 Sep 2024 19:57:53 +0000
Subject: [PATCH] Typo

---
 pdelfin/train/train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pdelfin/train/train.py b/pdelfin/train/train.py
index c803191..a1b4e2f 100644
--- a/pdelfin/train/train.py
+++ b/pdelfin/train/train.py
@@ -198,8 +198,8 @@ def run_train(config: TrainConfig):
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=formatted_dataset["train"],
-            eval_dataset=formatted_dataset["validation"],  # pyright: ignore
+            train_dataset=train_ds,
+            eval_dataset=validation_ds,
             tokenizer=processor.tokenizer,
             #Collator is not needed as we are doing batch size 1 for now...
             #data_collator=collator,