def preprocess(x, y):
    return x.view(-1, 1, 28, 28), y


class WrappedDataLoader:
    def __init__(self, dl, func):
        self.dl = dl
        self.func = func

    def __len__(self):
        return len(self.dl)

    def __iter__(self):
        for b in self.dl:
            yield (self.func(*b))

train_dl, valid_dl = get_data([train_ds](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.TensorDataset "torch.utils.data.TensorDataset"), [valid_ds](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.TensorDataset "torch.utils.data.TensorDataset"), bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)