diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index d785030..9d83ba9 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -41,7 +41,7 @@ class UpscaleDataset(Dataset):
         combined_df = pd.DataFrame()
         for parquet_file in parquet_files:
             # Load data with head() to limit rows for memory efficiency.
-            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(1250)
+            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(500)
             combined_df = pd.concat([combined_df, df], ignore_index=True)
         
         # Validate that each row has proper image formats.