onnx_tensorrt2.py

import numpy as np
import torch
from torch import nn
import onnx_tensorrt.backend as backend
import onnx

class Model(nn.Module):

    def forward(self, x):
        y = (2 * x)[0:1]
        return y

model = Model().eval()
dummy_input = torch.randn(4, 4, 4)

with torch.no_grad():
    torch.onnx.export(model, dummy_input, 'test.onnx', verbose=True)

model = onnx.load('test.onnx')
engine = backend.prepare(model, device='CUDA:0')