diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 3e7687a..1fefbe4 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -397,8 +397,8 @@ if __name__ == "__main__":
     trainer = FineTuner(
         model=model,
         dataset_paths=[
-            "/root/training_data/vision-dataset/image_upscaler.0",
-            "/root/training_data/vision-dataset/image_vec_upscaler.0"
+            "/root/training_data/vision-dataset/image_upscaler.parquet",
+            "/root/training_data/vision-dataset/image_vec_upscaler.parquet"
         ],
         batch_size=8,  # Increased batch size
         learning_rate=1e-4  # Reduced initial LR