Python set ONNX runtime to return tensor instead of numpy array

Question

In python I'm loading my predefined model (super-gradients, yolox-s):

onnx_session = onnxrt.InferenceSession("yolox_s_640_640.onnx")

Then I load some data and run it:

dataset = MyCostumeDataset(args.path, 'val')
val_dataloader = DataLoader(dataset, batch_size=args.bsize)

for inputs in val_dataloader:
    onnx_inputs = {onnx_session.get_inputs()[0].name: inputs}
    # inputs.shape: torch.Size([4, 3, 640, 640]), i.e., this is a Tensor
    
    raw_predictions = onnx_session.run(None, onnx_inputs)
    # this returns a list of numpy arrays:
    # type(raw_predictions[0])
    #  <class 'numpy.ndarray'>
    # raw_predictions[0].shape
    # (4, 8400, 85)

So far it is working as it should, except I'd like it to return, by default, a list of Tensors (torch.Tensor) instead of numpy array. I'm new to both ONNX and PyTorch, and I'm feeling like this is something basic that I'm missing here.

How can I get onnx_session to return a list of torch.Tensor, instead of numpy arrays? This will same some overhead in the conversion. Thanks!

Answer 1

According to the documentation you can get either - Returns: list of results, every result is either a numpy array, a sparse tensor, a list or a dictionary.

Documentation