ConfigDataset.hpp

#ifndef CONFIGDATASET__H
#define CONFIGDATASET__H

#include <torch/torch.h>
#include "Config.hpp"

class ConfigDataset : public torch::data::Dataset<ConfigDataset>
{
  private :

  std::vector<torch::Tensor> contexts;
  std::vector<torch::Tensor> classes;

  public :

  explicit ConfigDataset(std::vector<torch::Tensor> contexts, std::vector<torch::Tensor> classes);
  torch::optional<size_t> size() const override;
  torch::data::Example<> get(size_t index) override;
};

#endif