diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 46faf1a..6300266 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -13,7 +13,7 @@ class UpscaleDataset(Dataset):
         combined_df = pd.DataFrame()
         for parquet_file in parquet_files:
             # Load data with chunking for memory efficiency
-            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(5000)
+            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(4000)
             combined_df = pd.concat([combined_df, df], ignore_index=True)
         
         # Validate data format