From eafb094c6cf1660b6eb795ae58986294d5703822 Mon Sep 17 00:00:00 2001 From: Manfei <41607353+ManfeiBai@users.noreply.github.com> Date: Fri, 25 Oct 2024 00:11:28 -0700 Subject: [PATCH] [cherry-pick] R2.5.1 cherry pick #8315 (#8317) Co-authored-by: mcuiaws --- torch_xla/distributed/parallel_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_xla/distributed/parallel_loader.py b/torch_xla/distributed/parallel_loader.py index 7053361f795..22046df75e2 100644 --- a/torch_xla/distributed/parallel_loader.py +++ b/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]