if __name__ == '__main__':
    device = torch.device("cuda")
    model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
    model.to(device)
    model.eval()

    batch_size = 1
    image_shape = (3, 224, 224)
    data_shape = (batch_size,) + image_shape
    input_data = np.ones(data_shape).astype("float32")
    input_data = input_data * 10
    print(input_data)
    input_data = torch.from_numpy(input_data)
    device = torch.device("cuda")
    input_data = input_data.to(device)

    with torch.no_grad():
        output = model(input_data)
    torch.cuda.synchronize()

    # print(output[0])
    print(output[0][0:10]) // tensor([-5.8784, -1.7790, -0.6576, -4.9093, -3.7112, -1.5704, -7.2738,  2.6429,
                                       // -0.3956, -2.5655], device='cuda:0')

batch_size = 1
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape
out_shape = (batch_size, num_class)

mod, params = relay.testing.densenet.get_workload(
    densenet_size=169, batch_size=batch_size, image_shape=image_shape
)
opt_level = 3
target = tvm.target.cuda()

with tvm.transform.PassContext(opt_level=opt_level):
    lib = relay.build(mod, target, params=params)
    # graph_json, lib, params = relay.build(mod, target, params=params)  // if I use this style, the code that immediately follows to get the module will report an error, indicating that the module has no function 'default'

graph_json = lib.graph_json
params = lib.get_params()
ctx = tvm.gpu()
module = graph_runtime.GraphModule(lib["default"](ctx))

data = np.ones(data_shape).astype("float32")
data = data * 10
module.set_input("data", data)

module.run()

out = module.get_output(0, tvm.nd.empty(out_shape)).asnumpy()
print(out.flatten()[0:10]) //  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]