#ifndef CONFIGDATASET__H
#define CONFIGDATASET__H

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

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

  std::vector<std::unique_ptr<Config>> const & configs;
  std::vector<std::size_t> const & classes;
  std::size_t nbClasses;
  Dict & dict;

  public :

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

#endif