#include <torch/script.h>

torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU);

// device 0
torch::Device device(torch::kCUDA, 0);

int width = 512;
int height = 512;
int* pixelData = (int *) malloc(width * height * sizeof(int));
some_initialization(pixelData); // do something to load the data.
at::Tensor t = torch::from_blob(pixelData, { width, height });

auto options = c10::TensorOptions().dtype(torch::kShort);
at::Tensor t = torch::from_blob(pixelData, { 1, width, height }, options);

tensor = tensor.toType(torch::kFloat32);

// get the value of tensor on index(0, 0) to float
float val = tensor[0][0].item<float>()

class CustomDataset : public torch::data::Dataset<CustomDataset> {
  // use Batch as a alias of torch::data::Example<>
  using Batch = torch::data::Example<>;
  private:
    std::vector<std::string> file_list;

  public:
    explicit ICHDataset(const std::vector<std::string> file_list) 
      : file_list(file_list) {}

    Batch get(size_t index) {

      auto  = parse_data(index);

      return { data, label };
    } 

    torch::optional<size_t> size() const {
      return file_list.size();
    }
};

auto dataset = ICHDataset(file_list).map(torch::data::transforms::Stack<>());;;
auto dataloader = torch::data::make_data_loader(
  std::move(dataset),
  torch::data::DataLoaderOptions()
    .batch_size(batch_size)
    .workers(2)
    .enforce_ordering(true)
  );