TorchScriptModule.C

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

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

namespace Moose
{

TorchScriptModule::TorchScriptModule() {}

TorchScriptModule::TorchScriptModule(const std::string & filename) { loadNeuralNetwork(filename); }

void
TorchScriptModule::loadNeuralNetwork(const std::string & filename)
{
  try
  {
    torch::jit::script::Module * base = this;
    *base = torch::jit::load(filename);
  }
  catch (const c10::Error & e)
  {
    mooseError("Error while loading torchscript file ", filename, "!\n", e.msg());
  }
}

torch::Tensor
TorchScriptModule::forward(const torch::Tensor & x)
{
  std::vector<torch::jit::IValue> inputs(1, x);
  return torch::jit::script::Module::forward(inputs).toTensor();
}

}

#endif