[cherry-pick] R2.5.1 cherry pick #8315 (#8317)

Co-authored-by: mcuiaws <mikecui@amazon.com>

torch_xla/distributed/parallel_loader.py

@@ -131,7 +131,7 @@ def per_device_loader(self, device):
     return PerDeviceLoader(self, torch.device(device))
 
   def per_device_samples(self):
-    return len(self._loader) // len(self._devices)
+    return len(self._cpu_loader) // len(self._devices)
 
   def next_item(self, device):
     dqueue = self._queues[device]