use_primary_cuda_context

inference_models.models.common.cuda.use_primary_cuda_context

use_primary_cuda_context(cuda_device)

Context manager for using a CUDA device's primary context.

Retains and activates the primary CUDA context for a device, ensuring proper cleanup when exiting the context. This is useful when working with TensorRT engines or other CUDA operations that require an active context.

Parameters:

• cuda_device

  (Device) –

  PyCUDA Device object representing the CUDA device.

Yields:

• Context –

  cuda.Context: The active CUDA context for the device.

Examples:

Use primary CUDA context for TensorRT inference:

>>> from inference_models.developer_tools import use_primary_cuda_context
>>> import pycuda.driver as cuda
>>>
>>> cuda.init()
>>> device = cuda.Device(0)  # GPU 0
>>>
>>> with use_primary_cuda_context(device) as ctx:
...     # Perform CUDA operations
...     # TensorRT engine execution, etc.
...     pass

Note

• Requires PyCUDA to be installed
• Automatically pushes and pops the context
• Use this for TensorRT models or custom CUDA operations

See Also

• use_cuda_context(): Use an existing CUDA context