Using the torch::data::make_data_loader function in LibTorch

I’m having a bit of difficulty in LibTorch actually declaring and using torch::data::make_data_loader function call and then iterating over a dataset.

Supposing I have a header file sample_dataloader.h with code:

#include <torch/torch.h>

class sample_dataloader : public torch::data::Dataset<sample_dataloader> {
     public: 
     torch::Tensor rand_val;
     sample_dataloader(){
         rand_val = torch::rand({10, 3, 600, 600}); // random batch of images
     }
     
     torch::optional<size_t> size() const override {
         return 10;
     }
     
     torch::data::Example<> get(size_t index) {
         return rand_val[index];
     }
}

Then in my main.cpp file I try to use it as in:

#include <torch/torch.h>
#include "sample_dataloader.h"

int main() {
    auto train_set = sample_dataloader();
    int batch_size = 5;
    int num_workers = 2;
    auto dl = 
    torch::data::make_data_loader<torch::data::samplers::SequentialSampler> 
     (std::move(train_set), DataLoaderOptions().batch_size(batch_size).workers(num_workers));
    for (auto& batch : loader) {
       // do stuff 
    }
}

I get a compilation error that says “the range-based for statement requires a suitable begin function, but none was found”. When I try

for (auto& batch : *loader) {
       // do stuff 
}

instead, I get a memory error.

How am I supposed to use the torch::data::make_data_loader function properly?

Solved this problem, correct usage is:
“sample_dataloader.h”

#pragma once
#include <torch/torch.h>

class sample_dataloader : public torch::data::Dataset<sample_dataloader> {
public:
    torch::Tensor rand_val;
    sample_dataloader() {
        rand_val = torch::rand({ 10, 3, 600, 600 }); // random batch of images
    }

    torch::optional<size_t> size() const override {
        return 10;
    }

    torch::data::Example<> get(size_t index) {
        return { rand_val[index], torch::randint(5, {5}) };
    }
};

“sample_main.cpp”

#include <torch/torch.h>
#include "sample_dataloader.h"

int main() {
    auto train_set = sample_dataloader().map(torch::data::transforms::Stack<>());
    int batch_size = 5;
    int num_workers = 0;
    auto dl =
        torch::data::make_data_loader<torch::data::samplers::SequentialSampler>
        (std::move(train_set), torch::data::DataLoaderOptions().batch_size(batch_size).workers(num_workers));
    for (auto& batch : *dl) {
        auto rand_im = batch.data;
        auto targets = batch.target;
    }
}