Hi,
just in case anybody else is a little bit lost:
at::Tensor flatter(at::Tensor tensor_){
    return tensor_[0].flatten();
}

auto data_loader = torch::data::make_data_loader(
            torch::data::datasets::MNIST(PATH).map(
                    torch::data::transforms::TensorLambda<>(flatter)).map(
                    torch::data::transforms::Stack<>()), 64);