Skip to content

align_device_with_onnx_session

inference_models.models.common.onnx.align_device_with_onnx_session 

align_device_with_onnx_session(session, device, resolution_mode='fallback', fallback_device=None)

Make sure the declared torch device is in line with onnxruntime session capacity.

An onnxruntime session can only consume GPU-resident tensors when it runs a GPU-capable execution provider. When a model is initialized with a CUDA torch device (e.g. the auto-selected default on a GPU machine) but the session ended up CPU-only - either because the caller requested only CPUExecutionProvider or because onnxruntime silently fell back during initialization - binding CUDA tensors would fail at runtime with a cryptic "no data transfer registered" error. This function detects that mismatch upfront and resolves it according to resolution_mode, so pre- and post-processing stay on a device the session can read from.

Limitations

For now, only CUDA devices passed as the primary device are verified - any non-CUDA device is returned unchanged without validation.

Parameters:

  • session

    (InferenceSession) –

    Initialized ONNX Runtime session. Its effective (post-fallback) providers are read via session.get_providers().

  • device

    (device) –

    Torch device requested for the model.

  • resolution_mode

    (DeviceMismatchResolutionMode, default: 'fallback' ) –

    How to resolve a detected mismatch. "fallback" (default) logs a warning and returns the fallback device, "fail" raises EnvironmentConfigurationError.

  • fallback_device

    (Optional[device], default: None ) –

    Device to return when a mismatch is resolved in "fallback" mode. Default value (None) means default behaviour - falling back to the CPU device.

Returns:

  • device –

    The requested device when it is compatible with the session's providers,

  • device –

    otherwise the fallback device (in "fallback" mode).

Raises:

  • ModelInputError –

    When resolution_mode is not one of "fail", "fallback".

  • EnvironmentConfigurationError –

    When a mismatch is detected and resolution_mode is "fail".

Examples:

Align device in a model's from_pretrained:

>>> session = onnxruntime.InferenceSession(
...     "model.onnx", providers=["CPUExecutionProvider"]
... )
>>> device = align_device_with_onnx_session(
...     session=session, device=torch.device("cuda:0")
... )
>>> device
device(type='cpu')

Raise instead of falling back:

>>> device = align_device_with_onnx_session(
...     session=session,
...     device=torch.device("cuda:0"),
...     resolution_mode="fail",
... )
Traceback (most recent call last):
EnvironmentConfigurationError: ...