Skip to content
Snippets Groups Projects
ConfigDataset.cpp 364 B
Newer Older
#include "ConfigDataset.hpp"

Franck Dary's avatar
Franck Dary committed
ConfigDataset::ConfigDataset(std::vector<torch::Tensor> contexts, std::vector<torch::Tensor> classes) : contexts(contexts), classes(classes)
{
}

torch::optional<size_t> ConfigDataset::size() const
{
Franck Dary's avatar
Franck Dary committed
  return contexts.size();
}

torch::data::Example<> ConfigDataset::get(size_t index)
{
Franck Dary's avatar
Franck Dary committed
  return {contexts[index], classes[index]};