Python将ONNX运行时设置为返回张量而不是numpy数组

在python中，我正在加载预定义的模型（超级渐变，yolox-s）：

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

然后我加载一些数据并运行它：

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)

到目前为止，它仍在正常工作，但我希望它默认返回Tensor（torch.Tensor）列表，而不是numpy数组。我对ONNX和PyTorch都是新手，我觉得这是我在这里缺少的基本内容。

如何让onnx_session返回torch.Tensor列表，而不是numpy数组？这将使转换中的一些开销相同。谢谢