Output mismatch of torch.ceil due to an extra torch.transpose node when running on GPU

Description

When adding an extra node of torch.transpose as output in this model:

New:

The output of torch.ceil is expected to be the same for the same input in this 2 graphs. However, it mismatched bewteew the 2 models.

Environment

TensorRT Version: 8.6.1.post1

NVIDIA GPU: RTX 1660

NVIDIA Driver Version: 525.147.05

CUDA Version: 12.0

CUDNN Version: 8.9.4.25

Operating System: Ubuntu 22.04.3 LTS (x86_64)

Python Version (if applicable): 3.10.12

Tensorflow Version (if applicable): 2.13.0

PyTorch Version (if applicable): 2.1.0+cu118

Relevant Files

Model link:
models.zip

Input data file:
input_data.zip

Steps To Reproduce

Script:

from dataclasses import dataclass
from numpy import testing
import numpy as np
import torch
import tensorrt as trt
import pycuda.driver as cuda
from pycuda.driver import DeviceAllocation
import pickle


@dataclass
class HostDeviceMem:
    host: np.ndarray
    device: DeviceAllocation


class ONNXClassifierWrapper():
    def __init__(self, engine):
        self.engine = engine

    def allocate_memory(self):
        engine = self.engine
        inputs = []
        outputs = []
        bindings = []
        stream = cuda.Stream()
        onames = []
        name2idx = {}
        for idx, binding in enumerate(engine):
            name2idx[binding] = idx
            size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
            dtype = trt.nptype(engine.get_binding_dtype(binding))
            # Allocate host and device buffers
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            # Append the device buffer to device bindings.
            bindings.append(int(device_mem))
            # Append to the appropriate list.
            if engine.binding_is_input(binding):
                inputs.append(HostDeviceMem(host_mem, device_mem))
            else:
                outputs.append(HostDeviceMem(host_mem, device_mem))
                onames.append(binding)
        return inputs, outputs, bindings, stream, onames, name2idx

    def predict(self, inputs):  # result gets copied into output
        (
            trt_inputs,
            trt_outputs,
            trt_bindings,
            stream,
            onames,
            name2idx,
        ) = self.allocate_memory()
        context = self.engine.create_execution_context()
        # print(name2idx)
        # print(inputs.keys())
        for iname in inputs:
            np.copyto(
                trt_inputs[name2idx[iname]].host,
                inputs[iname]
                .astype(trt.nptype(self.engine.get_binding_dtype(iname)))
                .ravel(),
            )

        [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in trt_inputs]
        context.execute_async_v2(bindings=trt_bindings, stream_handle=stream.handle)
        [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in trt_outputs]
        stream.synchronize()
        trt_outputs = [out.host for out in trt_outputs]
        return {
            n: v.reshape(self.engine.get_binding_shape(n))
            for n, v in zip(onames, trt_outputs)
        }


def convert_onnx_to_engine(onnx_filename):
    logger = trt.Logger(trt.Logger.WARNING)
    with trt.Builder(logger) as builder, builder.create_network(
            1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
            trt.OnnxParser(network, logger) as parser:
        # builder.max_workspace_size = max_workspace_size
        # builder.fp16_mode = fp16_mode
        # builder.max_batch_size = max_batch_size

        config = builder.create_builder_config()
        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30)

        with open(onnx_filename, 'rb') as model:
            if not parser.parse(model.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))

        engine_bytes = builder.build_serialized_network(network, config)

        return trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(
            engine_bytes
        )

import pycuda.autoinit
DEVICE='cuda'

p0 = torch.tensor(45534, dtype=torch.float16)
p1 = torch.tensor(True, dtype=torch.bool)

class Model0(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.v6_0 = p0
        self.v5_0 = p1

    def forward(self, *args):
        _args = args
        v6_0 = self.v6_0
        v5_0 = self.v5_0
        getitem = _args[0]
        getitem_1 = _args[1]
        where = torch.where(v5_0, v6_0, getitem)
        ceil = torch.ceil(where)
        add = torch.add(where, getitem_1)
        return (ceil, add)

model_0 = Model0()
output_names_0 = ['v4_0', 'v3_0']
input_dict_0 = pickle.load(open('0.pickle', 'rb'))
inputs_0 = tuple(torch.from_numpy(v).to(DEVICE) for _, v in input_dict_0.items())
torch.onnx.export(model_0, inputs_0, '0.onnx', verbose=False, input_names=['v7_0', 'v2_0'], output_names=output_names_0, opset_version=14, do_constant_folding=False)

class Model1(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.v6_0 = p0
        self.v5_0 = p1

    def forward(self, *args):
        _args = args
        v6_0 = self.v6_0
        v5_0 = self.v5_0
        getitem = _args[0]
        getitem_1 = _args[1]
        where = torch.where(v5_0, v6_0, getitem)
        ceil = torch.ceil(where)
        transpose = getitem_1.transpose(0, 1)
        add = torch.add(where, getitem_1)
        return (ceil, transpose, add)

model_1 = Model1()
output_names_1 = ['v4_0', 'v9_0', 'v3_0']
input_dict_1 = pickle.load(open('0.pickle', 'rb'))
inputs_1 = tuple(torch.from_numpy(v).to(DEVICE) for _, v in input_dict_1.items())
torch.onnx.export(model_1, inputs_1, '1.onnx', verbose=False, input_names=['v7_0', 'v2_0'], output_names=output_names_1, opset_version=14, do_constant_folding=False)

engine_0 = convert_onnx_to_engine('0.onnx')
wrapper_0 = ONNXClassifierWrapper(engine_0)
output_0 = wrapper_0.predict(input_dict_0)

engine_1 = convert_onnx_to_engine('1.onnx')
wrapper_1 = ONNXClassifierWrapper(engine_1)
output_1 = wrapper_1.predict(input_dict_1)
output_name_dict = {'v3_0': 'v3_0',  'v4_0': 'v4_0'}

print('=========================')
try:
    for tensor_name_0, tensor_name_1 in output_name_dict.items():
        testing.assert_allclose(output_0[tensor_name_0], output_1[tensor_name_1], rtol=1, err_msg=f'at {tensor_name_0}, {tensor_name_1}')
    print("tensorRT does not trigger assertion")
except AssertionError as e:
    print("tensorRT triggers assertion")
    print(e)
print('=========================')

Steps to repro:

1. Download the input data file and put it at the same dir of the script.
2. Run the script

Output assertion:

=========================
tensorRT triggers assertion

Not equal to tolerance rtol=1, atol=0
at v4_0, v4_0
x and y -inf location mismatch:
 x: array([4.554e+04, 1.000e+00, 0.000e+00, 0.000e+00, 1.000e+00, 1.000e+00,
       1.000e+00, 0.000e+00,      -inf, 1.000e+00, 1.000e+00, 0.000e+00,
       1.000e+00, 1.000e+00, 0.000e+00, 0.000e+00, 0.000e+00, 1.000e+00,...
 y: array([4.554e+04, 1.000e+00, 0.000e+00, 0.000e+00, 1.000e+00, 1.000e+00,
       1.000e+00, 0.000e+00,      -inf, 1.000e+00, 1.000e+00, 0.000e+00,
       1.000e+00, 1.000e+00, 0.000e+00, 0.000e+00, 0.000e+00, 1.000e+00,...
=========================

Have you tried the latest release?: No

Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt): Yes