diff --git a/finetune_demo/finetune_vision.py b/finetune_demo/finetune_vision.py
index 6b47e37..f6693a4 100644
--- a/finetune_demo/finetune_vision.py
+++ b/finetune_demo/finetune_vision.py
@@ -236,6 +236,8 @@ class DataManager(object):
             batched=batched,
             remove_columns=remove_columns,
             num_proc=self._num_proc,
+            batch_size=500,    # Dynamically adjust according to mem size
+            writer_batch_size=500
         )