diff --git a/deepspeed/pt/deepspeed_light.py b/deepspeed/pt/deepspeed_light.py
index 6c335754bf50..20d680796127 100755
--- a/deepspeed/pt/deepspeed_light.py
+++ b/deepspeed/pt/deepspeed_light.py
@@ -650,8 +650,8 @@ def deepspeed_io(self,
         data_parallel_world_size = None
         data_parallel_rank = None
         if self.mpu is not None:
-            data_parallel_world_size = mpu.get_data_parallel_world_size()
-            data_parallel_rank = mpu.get_data_parallel_rank()
+            data_parallel_world_size = self.mpu.get_data_parallel_world_size()
+            data_parallel_rank = self.mpu.get_data_parallel_rank()
 
         return DeepSpeedDataLoader(dataset=dataset,
                                    batch_size=batch_size,