Isn't this just an Iterator that we can pipe from the original DataLoader iterator? Usually sub-classing specializes behavior, it shouldn't modify it . Maybe this is more like a decorator or something?
https://github.com/PyTorchLightning/pytorch-lightning/blob/c614cf0c3d7a252bfd01ff8516f9f42f1c1b055a/pytorch_lightning/lite/wrappers.py#L128