LArDLHelper.h

Go to the documentation of this file.

#ifndef LAR_DL_HELPER_H
#define LAR_DL_HELPER_H 1

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

#include "Pandora/StatusCodes.h"

namespace lar_dl_content
{

class LArDLHelper
{
public:
    typedef torch::jit::script::Module TorchModel;
    typedef torch::Tensor TorchInput;
    typedef std::vector<torch::jit::IValue> TorchInputVector;
    typedef at::Tensor TorchOutput;

    static pandora::StatusCode LoadModel(const std::string &filename, TorchModel &model);

    static void InitialiseInput(const at::IntArrayRef dimensions, TorchInput &tensor);

    static void Forward(TorchModel &model, const TorchInputVector &input, TorchOutput &output);
};

} // namespace lar_dl_content

#endif // #ifndef LAR_DL_HELPER_H