diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index e9c13d6..36ad40a 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -375,7 +375,7 @@ if __name__ == "__main__":
             "/root/training_data/vision-dataset/image_upscaler.parquet",
             "/root/training_data/vision-dataset/image_vec_upscaler.parquet"
         ],
-        batch_size=8,   # Increased batch size
+        batch_size=2,   # Increased batch size
         learning_rate=1e-4   # Reduced initial LR
     )