LibtorchModel.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

#pragma once

#ifdef NEML2_ENABLED

// libtorch headers
#include <torch/script.h>

// neml2 headers
#include "neml2/models/Model.h"

// moose headers
#include "DataFileUtils.h"

namespace neml2
{

class LibtorchModel : public Model
{
public:
  static OptionSet expected_options();

  LibtorchModel(const OptionSet & options);

  virtual void to(const torch::TensorOptions & options) override;

  virtual void request_AD() override;

protected:
  virtual void set_value(bool out, bool dout_din, bool d2out_din2) override;

  // Input variable vector
  std::vector<const Variable<Scalar> *> _inputs;
  // Output vector
  std::vector<Variable<Scalar> *> _outputs;
  Moose::DataFileUtils::Path _file_path;
  std::unique_ptr<torch::jit::script::Module> _surrogate;
};

}

#endif