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<Config> const & configs;
  std::vector<std::size_t> const & classes;

  public :

  ConfigDataset(std::vector<Config> const & configs, std::vector<std::size_t> const & classes);
  torch::optional<size_t> size() const override;
  torch::data::Example<> get(size_t index) override;
};

#endif