def __iter__(self) -> Iterator:
parallel_loader = ParallelLoader(self._loader, [self._device],
**self._parallel_loader_kwargs)
return parallel_loader.per_device_loader(self._device)
def autocast_loader(self, loader: DataLoader):
"""Loader wrapper for the distributed mode."""
return ParallelLoader(loader,
[self.device]).per_device_loader(self.device)