PhysicsTools/PyTorch/src/ScriptedModuleLoad.cc - master - githubmirror/cmssw

PhysicsTools/PyTorch/src/ScriptedModuleLoad.cc

24 строки1 KB

Christine Zeh

Integration of PyTorch based portable (direct) ML inference with Alpaka-based heterogeneous core

21 ноя 2025, 14:01

21 ноя 2025, 14:01a7218f1

#include "PhysicsTools/PyTorch/interface/ScriptModuleLoad.h"

namespace cms::torch {

  // `torch::jit::load` wrapper to load a JIT exported TorchScript model.
  // In case of failure, a cms::Exception is thrown with context and error details.
  // `dev` optional device to load the model on. Async load is not supported. Use model.to(device, true) instead.
  // TODO: figure out how to fix linker error without inlining and adding pytorch-cuda dependency (move impl to src/)
  //       that brings CUDA headers in alpaka ROCm/HIP compiled products.
  ScriptedModule load(const std::string &model_path, std::optional<::torch::Device> dev) {
    ScriptedModule model;
    try {
      model = ::torch::jit::load(model_path, dev);
    } catch (const c10::Error &e) {
      cms::Exception ex("ModelLoadingError");
      ex << "Error loading the model from path: " << model_path << "\n"
         << "Details: " << e.what();
      ex.addContext("Calling cms::torch::load(const std::string&)");
      throw ex;
    }
    return model;
  }

}  // namespace cms::torch