convert model to onnxruntime

[vuhungtvt2018]

how do you convert my model to onnxruntime?
please help me to resolve this problem!

[bayueyaoshao]

from models.unet import UNet
import torch
use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
model = UNet().to(device)
model_path = r"D:\Work\doc\RIAD\best.pt"
export_path = r"D:\Work\doc\RIAD\RIAD_model_320.onnx"
check_point = torch.load(model_path, map_location=device)
model.load_state_dict(check_point['model'])
inputs = torch.randn(1, 3, 320, 320)

if name == 'main':
torch.onnx.export(
model,
inputs,
export_path,
opset_version=12,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes=None
)