diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 6ce67b4..0226db4 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -15,7 +15,7 @@ from aiia import AIIA, AIIAConfig, AIIABase, AIIABaseShared, AIIAmoe, AIIAchunke
 class aiuNNDataset(torch.utils.data.Dataset):
     def __init__(self, parquet_path):
         # Read the Parquet file
-        self.df = pd.read_parquet(parquet_path)
+        self.df = pd.read_parquet(parquet_path).head(2500)
 
         # Data augmentation pipeline without Resize as it's redundant
         self.augmentation = Compose([