LibtorchDataset.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#ifdef LIBTORCH_ENABLED

#pragma once

#include <torch/torch.h>
#include "MooseError.h"

namespace Moose
{

class LibtorchDataset : public torch::data::datasets::Dataset<LibtorchDataset>
{
public:
  LibtorchDataset(torch::Tensor dt, torch::Tensor rt) : _data_tensor(dt), _response_tensor(rt) {}

  torch::data::Example<> get(size_t index) override
  {
    mooseAssert(index < size(), "Index is out of range!");
    return {_data_tensor[index], _response_tensor[index]};
  }

  torch::optional<size_t> size() const override
  {
    mooseAssert(_response_tensor.sizes().size(), "The tensors are empty!");
    return _response_tensor.sizes()[0];
  }

private:
  torch::Tensor _data_tensor;
  torch::Tensor _response_tensor;
};

}

#endif