diff --git a/python/paddle/fluid/dataset.py b/python/paddle/fluid/dataset.py
index d63773223ddc0c155f26a656f19c4ba80f482632..11ac326e8f7c44c6664dbcab1a979bbbb9123d1c 100644
--- a/python/paddle/fluid/dataset.py
+++ b/python/paddle/fluid/dataset.py
@@ -240,12 +240,17 @@ class InMemoryDataset(DatasetBase):
         Args:
             fleet: fleet singleton. Default None.
         """
+        trainer_id = 0
         trainer_num = 1
+        fleet_send_batch_size = 80000
         if fleet is not None:
             fleet.fleet_instance.role_maker_._barrier_worker()
+            trainer_id = fleet.worker_index()
             trainer_num = fleet.worker_num()
         self.dataset.register_client2client_msg_handler()
+        self.dataset.set_trainer_id(trainer_id)
         self.dataset.set_trainer_num(trainer_num)
+        self.dataset.set_fleet_send_batch_size(fleet_send_batch_size)
         if fleet is not None:
             fleet.fleet_instance.role_maker_._barrier_worker()
         self.dataset.global_shuffle()