import onnxruntime as ort
import torch
import torch.nn as nn
import numpy as np


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(4, 2)

    def forward(self, x):
        return self.fc1(x)


torch_model = MyModel()
torch_input = torch.randn(1, 4)
onnx_model = torch.onnx.dynamo_export(torch_model, torch_input)

onnx_model.save("model.onnx")

sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
binding = sess.io_binding()

for n in sess.get_inputs():
    binding.bind_ortvalue_input(
            n.name,
            ort.OrtValue.ortvalue_from_numpy(
                np.zeros(shape=n.shape, dtype=np.float32)))

for n in sess.get_outputs():
    binding.bind_ortvalue_output(
            n.name,
            ort.OrtValue.ortvalue_from_numpy(
                np.zeros(shape=n.shape, dtype=np.float32),
                device_type='cuda'))

sess.run_with_iobinding(binding)