diff --git a/hft.py b/hft.py
index b7ee50a..f960d70 100644
--- a/hft.py
+++ b/hft.py
@@ -130,7 +130,7 @@ training_args = TrainingArguments(
 trainer = CustomTrainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_dataset["train"],
+    train_dataset=tokenized_dataset,
 )
 
 # Trening modelu