Hello, I find that the input and output of the transformed model are the same, the model is not working correctly, how can I solve this problem?

[SuperCarryFu]

from torch2trt import * import torch from torchvision.models.alexnet import alexnet

if name == 'main': model = alexnet(pretrained=True).cuda().eval().half() data = torch.randn((1, 3, 224, 224)).cuda().half() data1 = data.contiguous() print('Running torch2trt...') model_trt = torch2trt(model, [data1],fp16_mode=True, max_workspace_size=1 << 25)

print('Saving model...')
torch.save(model_trt.state_dict(), '.test_model.pth')

print('Loading model...')
model_trt_2 = TRTModule()
model_trt_2.load_state_dict(torch.load('.test_model.pth'))

assert (model_trt_2.engine is not None)
#a=model_trt_2(data1)
#b=model_trt(data1)
#c=model(data)
print(torch.max(torch.abs(model_trt_2(data) - model(data))))
print(torch.max(torch.abs(model_trt_2(data) - model_trt(data))))