TensorRT: 🐛 [Bug] Torch Tensor RT crash when trying to compile a script module on Windows (C++)

Bug Description

I can’t compile a script module with TorchTensorRT. This is my code:

#include <iostream>
#include <vector>
#include <ATen/Context.h>
#include <torch/torch.h>
#include <torch/script.h>
#include "torch_tensorrt/torch_tensorrt.h"

void compile(std::string model_path) {

    const torch::Device device = torch::Device(torch::kCUDA, 0);
    torch::jit::script::Module model;

    std::cout << "Trying to load the model" << std::endl;
    try {
        model = torch::jit::load(model_path, device);
        model.to(device);
        model.eval();
        std::cout << "AI model loaded successfully." << std::endl;
    }
    catch (const c10::Error& e) {
        std::cerr << e.what() << std::endl;
    }

    auto input = torch_tensorrt::Input(std::vector<int64_t>{ 1, 3, 512, 512 });
    std::cout << "Creating compile settings" << std::endl;
    auto compile_settings = torch_tensorrt::ts::CompileSpec({ input });
    // Compile module
    std::cout << "Compiling..." << std::endl;
    auto trt_mod = torch_tensorrt::ts::compile(model, compile_settings);  <-- CRASHES HERE.
    // Run like normal
    std::cout << "Create tensor" << std::endl;
    auto in = torch::randn({ 1, 3, 512, 512 }, device);
    std::cout << "Forward pass..." << std::endl;
    auto results = trt_mod.forward({ in });
    // Save module for later
    trt_mod.save("output/model/path.ts");

}

int main() {

    compile("path/to/traced_script_module.pt");

    return 0;
}

This is the error I get:

First a WARNING get printed “WARNING: [Torch-TensorRT] - Interpolation layer will be run through ATen, not TensorRT. Performance may be lower than expected”, and then, as you can see from the screenshot, I got an exception “read access violation. creator was nullptr.” when running the following lines:

auto creator = getPluginRegistry()->getPluginCreator("Interpolate", "1", "torch_tensorrt");
auto interpolate_plugin = creator->createPlugin(name, &fc);

The file interpolate.cpp is located at path/to/Torch-TensorRT/core/conversion/converters/impl. What am I doing wrong?

This is my CMakeLists.txt:

cmake_minimum_required (VERSION 3.8)

project(example-app)

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

add_executable(example-app create_trt_module.cpp)

target_include_directories(example-app PRIVATE "path/to/Torch-TensorRT/cpp/include")

target_link_libraries(example-app "${TORCH_LIBRARIES}")
target_link_libraries(example-app  path/to/Torch-TensorRT/out/build/x64-Release/lib/torchtrt.lib) 
target_link_libraries(example-app  path/to/Torch-TensorRT/out/build/x64-Release/lib/torchtrt_plugins.lib)

I exported the traced script module with the following code:

# Import model
# ...
model.to("cuda")
model.eval()

# Create dummy data for tracing and benchmarking purposes.
shape = (1, 3, 512, 512)
input_data = torch.randn(shape).to("cuda")

# Convert model to script module
print("Tracing PyTorch model...")
traced_script_module = torch.jit.trace(model, input_data)
torch.jit.save(traced_script_module, "traced_script_module.pt")

Environment

• Torch-TensorRT Version: built from source from this branch (that is currently being merged) #1058
• TensorRT Version: 8.4.1.5
• CUDNN: 8.3.1
• CPU Architecture: x86-64
• OS : Windows 11
• LIbtorch: 1.11.0
• CUDA version: 11.5.2
• GPU model: NVIDIA RTX 3080 Mobile