peal.architectures.onnx_predictor.load_onnx_runtime¶

peal.architectures.onnx_predictor.load_onnx_runtime(path, device='cuda')[source]¶

The pre-conversion behaviour: an onnxruntime closure (not trainable).

Parameters:

path (str) – Path of the .onnx file.
device (str or torch.device) – Device the output tensor is moved to.

Returns:

onnx_model(input_data) that feeds the tensor to onnxruntime as float32 and returns the first output as a tensor on device. The closure carries the attribute onnx_path.

Return type:

callable