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