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


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)