如何测试pth和onnx模型推理结果是否一致

import torch
import onnxruntime as ort

# 加载pth模型
pth_model = torch.load("pth_model.pt")
pth_model.eval()

# 加载onnx模型
onnx_model = ort.InferenceSession("onnx_model.onnx")

# 加载测试图片
image = torch.randn(1, 3, 224, 224)

# 使用pth模型进行推理
with torch.no_grad():
    pth_output = pth_model(image)

# 使用onnx模型进行推理
onnx_output = onnx_model.run(None, {"input": image.cpu().numpy()})[0]
onnx_output = torch.from_numpy(onnx_output)

# 比较两个模型的推理结果
if torch.allclose(pth_output, onnx_output, rtol=1e-03, atol=1e-05):
    print("两个模型的推理结果相同")
else:
    print("两个模型的推理结果不同")