torch2trt
torch2trt copied to clipboard
Inconsistent inference results with avg_pool3d operator
Description:
I'm experiencing a discrepancy between the inference results of my PyTorch model and the TensorRT model obtained by converting it using the torch2trt tool.
Reproduce
This can be reproduced by the following script:
from torch2trt import torch2trt
import torch
from torch.nn import Module
para_0 = torch.randn([2, 3, 2, 4, 8], dtype=torch.float32)
para_1 = (2, 4, 8)
para_2 = 1
para_3 = (1, 1, 2)
para_4 = False
para_5 = True
para_6 = 1
class avg_pool3d(Module):
def forward(self, input):
return torch.nn.functional.avg_pool3d(input, para_1,para_2,para_3,para_4,para_5,para_6,)
model = avg_pool3d().float().cuda()
input_data=para_0.cuda()
model_trt = torch2trt(model, [input_data])
y = model(input_data)
y_trt = model_trt(input_data)
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
Environment
- torch: 1.11.0
- torch2trt: 0.4.0
- tensorrt: 8.6.1.6