peal.architectures.onnx_predictor.load_onnx_predictor¶

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

What get_predictor calls for a .onnx path: the trainable module when the conversion succeeds, the runtime closure otherwise.

Parameters:

path (str) – Path of the .onnx file.
device (str or torch.device) – Device for the module or the closure’s outputs.

Returns:

load_onnx_as_torch() result, or the load_onnx_runtime() closure when PEAL_ONNX_CONVERT=0 or the conversion raised ImportError/RuntimeError (a warning names the reason).

Return type:

OnnxPredictor or callable