TensorRT engine hangs when calling get_binding_shape with PyTorch imported

[zgspose]

Description:
I encountered a problem when trying to visualize inference results of ONNX and TensorRT (TRT) face detection models on the same image. The program hangs and terminates when initializing the TensorRT engine.

Steps to reproduce:

1. Load a TRT engine (face_detect_fp32.trt) using a shared CUDA context (shared_ctx).
2. Import PyTorch (import torch) before loading the engine.
3. In the TRTInference class, during initialization, the line:

shape = tuple(self.engine.get_binding_shape(i))

causes the program to hang.

4. If I do not import PyTorch, everything works normally.

Code snippet (simplified):

import torch  # <- importing this causes the hang
import pycuda.driver as cuda
import tensorrt as trt

cuda.init()
ctx = cuda.Device(0).retain_primary_context()
ctx.push()

class TRTInference:
    def __init__(self, engine_path: str, ctx=None):
        with open(engine_path, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        for i in range(self.engine.num_bindings):
            shape = tuple(self.engine.get_binding_shape(i))  # <--- hangs here

Observation:

• The hang occurs only when importing PyTorch.
• Using a shared CUDA context for TensorRT (ctx = cuda.Device(0).retain_primary_context()) does not prevent the issue.
• I suspect there is a CUDA context conflict between PyTorch and PyCUDA/TensorRT.

Question:

• Why does importing PyTorch cause get_binding_shape() to hang?
• How can I modify the code to allow PyTorch and TensorRT to work in the same process safely?

Environment:

• Python 3.8
• PyTorch, TensorRT, PyCUDA
• Shared CUDA context

[zgspose]

I mainly use PyTorch for data preprocessing and postprocessing operations.