diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 03efb8f..0013ec9 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 a subset (head(2500)) from each parquet file
-            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(2500)
+            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(1250)
             combined_df = pd.concat([combined_df, df], ignore_index=True)
 
         # Validate rows (ensuring each value is bytes or str)