torch_tensorrt.dynamo.compile creates flawed output

[AIManifest]

Greetings friends,

I wanted to reach out for assistance with the compilation of my model. The issue I'm having is this: I'm successfully converting the model using torch_tensorrt.dynamo.compile, however, the output that I get is flawed. I'm modifying https://github.com/dajes/frame-interpolation-pytorch by converting the model to torch_tensorrt. I was wondering if any of you experts can chime in with helping me properly compile the model? I must be missing something because the model converts correctly but it just doesn't produce output as expected. This is important to me because this method of compilation produces 2x faster results than torch.compile and torch_tensorrt.compile. Any input and assistance would be greatly appreciated. Please see below for the code:

import os
import gc
import cv2
import math
import time
import torch
import bisect
import warnings
import tensorrt
import threading
import tensorflow
import skvideo.io
import torch_tensorrt
import numpy as np
import tensorflow as tf
from glob import glob
from queue import Queue, Empty
from .film_util import load_image
from types import SimpleNamespace
from tqdm.auto import tqdm, trange
from .interpolator import Interpolator
from dataclasses import dataclass

warnings.filterwarnings("ignore")

class Backend:
    @dataclass
    class Eager:
        module: torch.nn.Module

    @dataclass
    class TensorRT:
        module: list[torch.nn.Module]

def clear_memory():
    gc.collect()
    print("Emptied RAM")
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    torch.cuda.synchronize()
    print("Emptied VRAM")

def verbose_print(string, do_print=True):
    if do_print:
        print(string)

def translate_state_dict(var_dict, state_dict):
    for name, (prev_name, weight) in zip(state_dict, var_dict.items()):
        verbose_print(f'Mapping, {prev_name}, ->, {name}', do_print=False)
        weight = torch.from_numpy(weight)
        if 'kernel' in prev_name:
            # Transpose the conv2d kernel weights, since TF uses (H, W, C, K) and PyTorch uses (K, C, H, W)
            weight = weight.permute(3, 2, 0, 1)

        assert state_dict[name].shape == weight.shape, f'Shape mismatch {state_dict[name].shape} != {weight.shape}'

        state_dict[name] = weight


def import_state_dict(interpolator: Interpolator, saved_model):
    variables = saved_model.keras_api.variables

    extract_dict = interpolator.extract.state_dict()
    flow_dict = interpolator.predict_flow.state_dict()
    fuse_dict = interpolator.fuse.state_dict()

    extract_vars = {}
    _flow_vars = {}
    _fuse_vars = {}

    for var in variables:
        name = var.name
        if name.startswith('feat_net'):
            extract_vars[name[9:]] = var.numpy()
        elif name.startswith('predict_flow'):
            _flow_vars[name[13:]] = var.numpy()
        elif name.startswith('fusion'):
            _fuse_vars[name[7:]] = var.numpy()

    # reverse order of modules to allow jit export
    # TODO: improve this hack
    flow_vars = dict(sorted(_flow_vars.items(), key=lambda x: x[0].split('/')[0], reverse=True))
    fuse_vars = dict(sorted(_fuse_vars.items(), key=lambda x: int((x[0].split('/')[0].split('_')[1:] or [0])[0]) // 3, reverse=True))

    assert len(extract_vars) == len(extract_dict), f'{len(extract_vars)} != {len(extract_dict)}'
    assert len(flow_vars) == len(flow_dict), f'{len(flow_vars)} != {len(flow_dict)}'
    assert len(fuse_vars) == len(fuse_dict), f'{len(fuse_vars)} != {len(fuse_dict)}'

    for state_dict, var_dict in ((extract_dict, extract_vars), (flow_dict, flow_vars), (fuse_dict, fuse_vars)):
        translate_state_dict(var_dict, state_dict)

    interpolator.extract.load_state_dict(extract_dict)
    interpolator.predict_flow.load_state_dict(flow_dict)
    interpolator.fuse.load_state_dict(fuse_dict)


def verify_debug_outputs(pt_outputs, tf_outputs):
    max_error = 0
    for name, predicted in pt_outputs.items():
        if name == 'image':
            continue
        pred_frfp = [f.permute(0, 2, 3, 1).detach().cpu().numpy() for f in predicted]
        true_frfp = [f.numpy() for f in tf_outputs[name]]

        for i, (pred, true) in enumerate(zip(pred_frfp, true_frfp)):
            assert pred.shape == true.shape, f'{name} {i} shape mismatch {pred.shape} != {true.shape}'
            error = np.max(np.abs(pred - true))
            max_error = max(max_error, error)
            assert error < 1, f'{name} {i} max error: {error}'
    print('Max intermediate error:', max_error)


def test_model(interpolator, model, half=False, gpu=False):
    torch.manual_seed(0)
    time = torch.full((1, 1), .5)
    x0 = torch.rand(1, 3, 256, 256)
    x1 = torch.rand(1, 3, 256, 256)

    x0_ = tf.convert_to_tensor(x0.permute(0, 2, 3, 1).numpy(), dtype=tf.float32)
    x1_ = tf.convert_to_tensor(x1.permute(0, 2, 3, 1).numpy(), dtype=tf.float32)
    time_ = tf.convert_to_tensor(time.numpy(), dtype=tf.float32)
    tf_outputs = model({'x0': x0_, 'x1': x1_, 'time': time_}, training=False)

    if half:
        x0 = x0.half()
        x1 = x1.half()
        time = time.half()

    if gpu and torch.cuda.is_available():
        x0 = x0.cuda()
        x1 = x1.cuda()
        time = time.cuda()

    with torch.no_grad():
        pt_outputs = interpolator.debug_forward(x0, x1, time)

    verify_debug_outputs(pt_outputs, tf_outputs)

    with torch.no_grad():
        prediction = interpolator(x0, x1, time)
    output_color = prediction.permute(0, 2, 3, 1).detach().cpu().numpy()
    true_color = tf_outputs['image'].numpy()
    error = np.abs(output_color - true_color).max()

    print('Color max error:', error)


def create_model(model_path, save_path, trt=False,export_to_torchscript=True, use_gpu=False, 
                 fp16=True, skiptest=False, video_path=None):
    print(f'Exporting model to FP{["32", "16"][fp16]} {["state_dict", "torchscript"][export_to_torchscript]} '
          f'using {"CG"[use_gpu]}PU')
    model = tf.compat.v2.saved_model.load(model_path)
    interpolator = Interpolator()
    interpolator.eval().cuda()
    import_state_dict(interpolator, model)
    interpolator.eval().cuda()
    import_state_dict(interpolator, model)
    if trt:
        if not os.path.exists(save_path) and not export_to_torchscript:
            torch.save(interpolator.state_dict(), save_path)
        # else:
            # scripted_model = torch.jit.script(interpolator)
            # scripted_model.save(save_path)
        if not export_to_torchscript:
            interpolator_sd = torch.load(save_path, weights_only=True, map_location="cpu")
            interpolator.load_state_dict(interpolator_sd)
    interpolator_model = interpolator
    if fp16: 
        print("Running FP16 Mode")
        interpolator_model.half()

    print("success trt")
    path = video_path
    cap = cv2.VideoCapture(path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    trt_base_shape = [height, width]
    cap.release()
    trt_min_shape = [64, 64]
    trt_opt_shape = [width, height]
    trt_max_shape = [width*4, height*4]
    trt_cache_dir = os.path.dirname(model_path)
    fp16 = fp16
    trt_workspace_size = 0
    trt_max_aux_streams = 2
    trt_optimization_level = 5
    trt_debug = False
    device_index = 0
    device = torch.device("cuda", device_index)
    dtype = torch.half if fp16 else torch.float32
    require_full_compilation = True
    truncate_double = True
    require_full_compilation = True
    
    w = width
    h = height

    tmp = max(64, int(64 / 1.0))
    # pw = math.ceil(w / tmp) * tmp
    # ph = math.ceil(h / tmp) * tmp
    # padding = (0, pw - w, 0, ph - h)
    # tmp = max(128, int(128 / scale))
    ph = ((h - 1) // tmp + 1) * tmp
    pw = ((w - 1) // tmp + 1) * tmp
    padding = (0, pw - w, 0, ph - h)
    
    for i in range(2):
        trt_min_shape[i] = math.ceil(max(trt_min_shape[i], 1) / tmp) * tmp
        trt_opt_shape[i] = math.ceil(max(trt_opt_shape[i], 1) / tmp) * tmp
        trt_max_shape[i] = math.ceil(max(trt_max_shape[i], 1) / tmp) * tmp
    
    dimensions = (
        f"min-{trt_min_shape[0]}x{trt_min_shape[1]}"
        f"_opt-{trt_opt_shape[0]}x{trt_opt_shape[1]}"
        f"_max-{trt_max_shape[0]}x{trt_max_shape[1]}"
    )
    
    trt_engine_path = os.path.join(
        os.path.realpath(trt_cache_dir),
        (
            f"film"
            + f"_{dimensions}"
            + f"_{'fp16' if fp16 else 'fp32'}"
            + f"_{torch.cuda.get_device_name(device)}"
            + f"_trt-{tensorrt.__version__}"
            + (f"_workspace-{trt_workspace_size}" if trt_workspace_size > 0 else "")
            + (f"_aux-{trt_max_aux_streams}" if trt_max_aux_streams is not None else "")
            + (f"_level-{trt_optimization_level}" if trt_optimization_level is not None else "")
            + ".ts"
        ),
    )
    print(f"TRT ENGINE OUTPUT PATH: {trt_engine_path}")
    
    if not os.path.isfile(trt_engine_path) and trt:
        trt_min_shape.reverse()
        trt_opt_shape.reverse()
        trt_max_shape.reverse()

        x0 = torch.zeros((1, 3, ph, pw), dtype=dtype, device=device)
        x1 = torch.zeros((1, 3, ph, pw), dtype=dtype, device=device)
        batch_dt = x0.new_full((1, 1), .5)

        example_tensors = (
            x0,
            x1,
            batch_dt
        )

        example_tensors_list = [
            x0,
            x1,
            batch_dt
        ]
    
        _height = torch.export.Dim("height", min=trt_min_shape[0] // tmp, max=trt_max_shape[0] // tmp)
        _width = torch.export.Dim("width", min=trt_min_shape[1] // tmp, max=trt_max_shape[1] // tmp)
        _dt = torch.export.Dim("timestep")
        dim_height = _height * tmp
        dim_width = _width * tmp

        dynamic_shapes = {
            "x0": {2: dim_height, 3: dim_width},
            "x1": {2: dim_height, 3: dim_width},
            "batch_dt": {0: _dt, 1: _dt}
        }
    
        exported_program = torch.export.export(interpolator_model, example_tensors)#, dynamic_shapes=dynamic_shapes)

        clear_memory()
        
        trt_dtype = torch_tensorrt.dtype.float16 if fp16 else torch_tensorrt.dtype.float32
        
        inputs = [
            torch_tensorrt.Input(
                shape=[1, 3] + trt_opt_shape,
                # min_shape=[1, 3] + trt_min_shape,
                # opt_shape=[1, 3] + trt_opt_shape,
                # max_shape=[1, 3] + trt_max_shape,
                device=device,
                dtype=trt_dtype,
                name="x0",
            ),
            torch_tensorrt.Input(
                shape=[1, 3] + trt_opt_shape,
                # min_shape=[1, 3] + trt_min_shape,
                # opt_shape=[1, 3] + trt_opt_shape,
                # max_shape=[1, 3] + trt_max_shape,
                device=device,
                dtype=trt_dtype,
                name="x1",
            ),
            torch_tensorrt.Input(
                shape=batch_dt.shape,
                # opt_shape=[1, 1],
                # max_shape=[1, 1],
                device=device,
                dtype=trt_dtype,
                name="batch_dt"
            )
        ]
        
        compiled_model = torch_tensorrt.dynamo.compile(
            exported_program,
            inputs,
            enabled_precisions={torch_tensorrt.dtype.float16, torch_tensorrt.dtype.float32},
            debug=trt_debug,
            workspace_size=trt_workspace_size,
            min_block_size=833,
            max_aux_streams=5,
            optimization_level=5,
            device=device,
            assume_dynamic_shape_support=True,
            require_full_compilation=False,
            truncate_double=truncate_double,
            dryrun=False,
            use_fast_partitioner=True,
            torch_executed_modules={
                "interpolator.extract",
                "interpolator.predict_flow",
                "interpolator.fuse"
            }
            # disable_tf32=fp16
        )

        print(f"Compiled Module Graph: {compiled_model.graph}")
        torch_tensorrt.save(compiled_model, trt_engine_path, output_format="torchscript", inputs=example_tensors)
        
    if not trt:
        if use_gpu and torch.cuda.is_available():
            interpolator = interpolator.cuda()
        else:
            use_gpu = False
    
        if fp16:
            interpolator = interpolator.half()
        if export_to_torchscript:
            interpolator = torch.jit.script(interpolator)
        if export_to_torchscript:
            interpolator.save(save_path)
        else:
            torch.save(interpolator.state_dict(), save_path)
    
        print("Returning Created FILM Model")
        # interpolator = torch.jit.load(save_path, map_location='cpu')
        if fp16:
            interpolator = interpolator.half()
        interpolator = interpolator.cuda().eval()
    else:
        if not os.path.exists(trt_engine_path):
            print("Returning Created TorchTensorRT Engine")
            interpolator = compiled_model.half().cuda().eval() if fp16 else compiled_model.cuda().eval()
        else:
            print("Loading Pre-Saved Created TorchTensorRT Engine")
            compiled_model = torch.jit.load(trt_engine_path).cuda().eval()
            interpolator = compiled_model.half() if fp16 else compiled_model
        # if os.path.exists(trt_engine_path):
        #     interpolator = compiled_model.cuda().eval()
        #     if fp16:
        #         interpolator = interpolator.half()
        # else:
        #     # Load the scripted module
        #     compiled_model = torch.jit.load(trt_engine_path)
            
        #     # Initialize the backend
        #     backend = Backend.TensorRT(module=compiled_model)
        #     interpolator = compiled_model.cuda().eval()
        # if fp16:
        #     interpolator = interpolator.half()

    return interpolator


@torch.inference_mode()
def run_film_interp_infer(
    model_path=None,
    input_folder=None,
    save_folder=None,
    inter_frames=None,
    video_path=None,
    trt=False,
    fp16=True,
    film_tensorflow_model_path=None):
    print(f" | PASSING THESE SETTINGS INTO FILM INTERPOLATOR: MODEL PATH: {model_path} | INPUT FOLDER: {input_folder} | SAVE FOLDER: {save_folder} | INTERPOLATING FRAMES {inter_frames} TIMES | INTERPOLATING FROM VIDEO PATH: {video_path} | USING TRT: {trt} | RUNNING IN FP16 MODE: {fp16} | BUILDING MODEL FROM: {film_tensorflow_model_path} | ")
    
    args = SimpleNamespace()
    args.model_path = model_path
    args.input_folder = input_folder
    args.save_folder = save_folder
    args.inter_frames = inter_frames
    args.video_path = video_path

    #DTYPE Settings
    dtype_name = "fp16" if fp16 else "fp32"

    if args.video_path is None:
        # Check if the folder exists
        if not os.path.exists(args.input_folder):
            print(f"Error: Folder '{args.input_folder}' does not exist.")
            return
        # Check if the folder contains any PNG or JPEG images
        if not any([f.endswith(".png") or f.endswith(".jpg") for f in os.listdir(args.input_folder)]):
            print(f"Error: Folder '{args.input_folder}' does not contain any PNG or JPEG images.")
            return
    
    # Sort Jpg/Png images by name
    if args.video_path is None:
        image_paths = sorted(glob(os.path.join(args.input_folder, "*.[jJ][pP][gG]")) + glob(os.path.join(args.input_folder, "*.[pP][nN][gG]")))
        print(f"Total frames to FILM-interpolate: {len(image_paths)}. Total frame-pairs: {len(image_paths)-1}.")
    else:
        #Video Ops
        cap = cv2.VideoCapture(args.video_path)
        videogen = skvideo.io.vreader(args.video_path)
        lastframe = next(videogen)
        #Frame Info
        height, width, channels = lastframe.shape
        #Video Info
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        #Interpolation Vars
        frames_to_add = (frame_count * inter_frames - frame_count) / (frame_count - 1)
        outfps_mul = int(round(frames_to_add))
        args.inter_frames, inter_frames = outfps_mul, outfps_mul
        output_fps = int(fps * (args.inter_frames + 1))
        #Output Video Settings
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        clip_name = os.path.splitext(args.video_path)[0]
        clip_dir = os.path.dirname(args.video_path)
        name_length = len([f for f in os.listdir(clip_dir) if os.path.isfile(os.path.join(clip_dir, f)) or clip_name in f])
        clip_output_path = os.path.join(clip_dir, f"{clip_name}_{name_length+1:03d}_{dtype_name}_{'trt_' if trt else ''}output_clip.mp4")
        writer = cv2.VideoWriter(clip_output_path, fourcc, output_fps, (width, height))
        print(f'| Running with INPUT FRAME HEIGHT: {height} | INPUT FRAME WIDTH {width} | INPUT FPS: {fps} | TOTAL ORIGINAL VIDEO FRAME COUNT: {frame_count}')
        print(f'| Running with OUTPUT FPS: {output_fps} | INTERPOLATING FRAMES |{inter_frames}| TIMES | Saving Video to: {clip_output_path} | ')

    #Torch_TensorRT (GODSend:))
    # export_to_torchscript=True, use_gpu=False, fp16=True, skiptest=False
    dirname = os.path.dirname(film_tensorflow_model_path)
    save_path = os.path.join(dirname, f"filmnet_{dtype_name}.pt")
    model = create_model(film_tensorflow_model_path, save_path, trt=trt, export_to_torchscript=True, use_gpu=True, fp16=fp16, skiptest=True, video_path=video_path)
    # if trt:
    #     print("Running Torch_TensoRT Dynamo Compile")
    #     # model = torch.compile(model, backend="torch_tensorrt")
    #     dirname = os.path.dirname(film_tensorflow_model_path)
    #     film_tensorflow_model_path = film_tensorflow_model_path
    #     save_path = os.path.join(dirname, f"filmnet_{dtype_name}.pt")
        
    #     fp32 = not fp16
    #     gpu = False
    #     skiptest = False
    #     statedict = False
    #     # export_to_torchscript=True, use_gpu=False, fp16=True, skiptest=False
    #     model = create_model(film_tensorflow_model_path, save_path, trt=trt, export_to_torchscript=statedict, use_gpu=True, fp16=fp16, skiptest=skiptest, video_path=video_path)
    # else:
    #     print("Loading Pytorch File")
    #     if os.path.exists(model_path):
    #         model = torch.jit.load(model_path, map_location='cpu')
    #     else:
    #         print("Creating PyTorch Model From Scratch")
    #         dirname = os.path.dirname(film_tensorflow_model_path)
    #         film_tensorflow_model_path = film_tensorflow_model_path
    #         save_path = os.path.join(dirname, f"filmnet_{dtype_name}.pt")
            
    #         fp32 = not fp16
    #         gpu = False
    #         skiptest = False
    #         statedict = True
    #         model = create_model(film_tensorflow_model_path, save_path, trt=trt, export_to_torchscript=statedict, use_gpu=True, fp16=fp16, skiptest=skiptest, video_path=video_path)
    #         del model
    #         model = torch.jit.load(model_path, map_location='cpu')
    
    # Select Precision Type
    if fp16:
        model = model.half()
    else:
        model = model.float()
    #Activate Model
    if not trt:
        model.eval()
    model.cuda()

    #Clear the Write Buffer Func
    def clear_write_buffer(write_buffer):
        while True:
            item = write_buffer.get()
            if item is None:
                break
            writer.write(item)
            
    #Build the Read Buffer Func
    def build_read_buffer(read_buffer, videogen):
        for frame in videogen:
            read_buffer.put(frame)
        read_buffer.put(None)

    #Initiate Buffers
    write_buffer = Queue(maxsize=500)
    read_buffer = Queue(maxsize=500)
    #Start Threads
    threading.Thread(target=build_read_buffer, args=(read_buffer, videogen)).start()
    threading.Thread(target=clear_write_buffer, args=(write_buffer,)).start()

    start_time = time.time()  # Timer START

    #Start First Frame Outside of the Loop Because Logic
    img1 = lastframe if input_folder is None else None

    #Range of Frames to be Interpolated
    film_range = len(image_paths) - 1 if video_path is None else frame_count - 1

    #Begin Interpolation Loop
    for i in trange(film_range):
        #Video Interpolation
        if args.video_path is not None:
            frame = read_buffer.get()
            if frame is None:
                break
            img2 = frame
        #Images Interpolation
        else:
            img1 = cv2.imread(image_paths[i])
            img2 = cv2.imread(image_paths[i + 1])
        
        img_batch_1, crop_region_1 = load_image(img1)
        img_batch_2, crop_region_2 = load_image(img2)
        img_batch_1 = torch.from_numpy(img_batch_1).permute(0, 3, 1, 2).float().cuda()
        img_batch_2 = torch.from_numpy(img_batch_2).permute(0, 3, 1, 2).float().cuda()

        results = [img_batch_1, img_batch_2]
        idxes = [0, args.inter_frames + 1]
        remains = list(range(1, args.inter_frames + 1))
        splits = torch.linspace(0, 1, args.inter_frames + 2)

        for _ in tqdm(range(len(remains)), leave=False, desc="In-Betweens Progress", disable=True):
            starts = splits[idxes[:-1]]
            ends = splits[idxes[1:]]
            distances = ((splits[None, remains] - starts[:, None]) / (ends[:, None] - starts[:, None]) - 0.5).abs()
            matrix = torch.argmin(distances).item()
            start_i, step = np.unravel_index(matrix, distances.shape)
            end_i = start_i + 1

            x0 = results[start_i]
            x1 = results[end_i]

            if fp16:
                x0 = x0.half()
                x1 = x1.half()
            else:
                x0 = x0.float()
                x1 = x1.float()
            x0 = x0.cuda()
            x1 = x1.cuda()

            dt = x0.new_full((1, 1), (splits[remains[step]] - splits[idxes[start_i]])) / (splits[idxes[end_i]] - splits[idxes[start_i]])

            with torch.no_grad():
                prediction = model(x0, x1, dt)
            insert_position = bisect.bisect_left(idxes, remains[step])
            idxes.insert(insert_position, remains[step])
            results.insert(insert_position, prediction.clamp(0, 1).cpu())
            del remains[step]

        y1, x1, y2, x2 = crop_region_1
        frames = [(tensor[0].cpu() * 255).byte().flip(0).permute(1, 2, 0).cpu().numpy()[y1:y2, x1:x2].copy() for tensor in results]
        
        for frame in frames:
            write_buffer.put(frame)

        img1 = img2
        
    write_buffer.put(None)
    
    while(not write_buffer.empty()):
        time.sleep(0.1)

    if args.video_path is not None:
        cap.release()
        writer.release()

    clear_memory()
    print(f"Interpolation done in {time.time() - start_time:.2f} seconds!")
    return clip_output_path

[AIManifest]

For anyone reading into this thread, I figured I need to show you exactly what I mean, so here are the comparisons:

Original Input:
https://drive.google.com/file/d/1yBSu5kU2llJbRLEnMAEgvVT6IkVQd6vt/view?usp=share_link

Output with Torch TensorRT:
https://drive.google.com/file/d/1a92iG4k1mdTx9w0rUukLhw8l7LmRYX-d/view?usp=share_link

Output without TRT behaving as expected:
https://drive.google.com/file/d/1UDEMAzn_s2L9kg7jR-aAacM6cj3YPBSF/view?usp=share_link

[AIManifest]

Here's a printed version of the outputs. They're completely different...

`pytorch_output`: tensor([[[[-7.0572e-04, -1.5640e-04, -3.4332e-05,  ..., -3.8147e-06,
           -1.6785e-04, -6.4087e-04],
          [ 2.4033e-04,  3.8147e-06,  1.9455e-04,  ..., -3.8147e-05,
            1.3733e-04,  1.7548e-04],
          [ 3.8147e-05,  1.1444e-05,  4.1962e-05,  ...,  8.3923e-05,
            1.9836e-04, -9.5367e-05],
          ...,
          [ 4.1962e-05,  4.1962e-05, -3.8147e-06,  ...,  5.3406e-05,
            4.9591e-05, -1.0300e-04],
          [-7.8201e-04,  1.2589e-04,  2.8992e-04,  ...,  3.0518e-05,
            6.1035e-05,  9.1553e-05],
          [ 2.0599e-04, -1.6785e-04,  9.1553e-05,  ...,  6.1035e-05,
            2.2888e-05, -4.5013e-04]],

         [[ 8.0109e-04, -9.1553e-05, -2.6703e-05,  ...,  1.2207e-04,
           -1.3351e-04, -5.2643e-04],
          [-3.8147e-05,  9.9182e-05,  1.2970e-04,  ...,  1.5259e-05,
            1.2589e-04, -4.5776e-05],
          [ 3.0518e-05,  1.0300e-04,  1.1444e-05,  ..., -2.2888e-05,
            1.1444e-05,  6.4850e-05],
          ...,
          [ 1.2207e-04,  3.8147e-05,  2.6703e-05,  ...,  4.5776e-05,
            5.3406e-05,  6.1035e-05],
          [-4.5013e-04,  1.0681e-04,  1.3733e-04,  ...,  2.2888e-05,
            1.2589e-04,  1.4496e-04],
          [-5.6839e-04,  2.5940e-04,  3.0518e-05,  ..., -1.9455e-04,
           -2.8610e-04,  1.6022e-04]],

         [[ 6.6757e-05,  1.0872e-04, -1.1158e-04,  ..., -9.1553e-05,
            4.6730e-05,  1.0920e-03],
          [ 3.3379e-05,  3.4046e-04,  5.7220e-05,  ...,  8.8692e-05,
            1.5926e-04,  2.1935e-04],
          [-2.8610e-06,  1.0109e-04,  0.0000e+00,  ..., -2.6703e-05,
            3.5286e-05,  8.5831e-05],
          ...,
          [-1.2350e-04,  2.4700e-04, -1.9073e-06,  ...,  1.2875e-04,
            1.9741e-04,  6.8665e-05],
          [-4.1437e-04, -1.2398e-05, -1.0014e-04,  ..., -1.2398e-05,
           -2.8610e-05,  2.4223e-04],
          [ 7.1335e-04,  7.8201e-05, -1.3733e-04,  ...,  1.5831e-04,
            1.1063e-04,  6.4182e-04]]]], device='cuda:0', dtype=torch.float16,
       grad_fn=<ConvolutionBackward0>)
`trt_output`: tensor([[[[ 3.8635e-02,  3.6194e-02,  3.4851e-02,  ...,  3.5034e-02,
            3.7170e-02,  3.4912e-02],
          [ 3.4729e-02,  3.2135e-02,  3.1250e-02,  ...,  3.1403e-02,
            3.3508e-02,  3.3600e-02],
          [ 3.4393e-02,  3.2135e-02,  3.0716e-02,  ...,  3.0701e-02,
            3.2104e-02,  3.1860e-02],
          ...,
          [ 3.4546e-02,  3.1235e-02,  3.0045e-02,  ...,  3.0365e-02,
            3.1799e-02,  3.2257e-02],
          [ 3.4943e-02,  3.1982e-02,  3.1494e-02,  ...,  3.1311e-02,
            3.1067e-02,  3.3752e-02],
          [ 3.3691e-02,  3.3844e-02,  3.3112e-02,  ...,  3.2745e-02,
            3.1052e-02,  3.5339e-02]],

         [[-7.3242e-02, -7.6660e-02, -7.4585e-02,  ..., -7.4097e-02,
           -7.5012e-02, -7.3059e-02],
          [-7.1777e-02, -7.3792e-02, -7.0618e-02,  ..., -7.0312e-02,
           -7.1045e-02, -6.9092e-02],
          [-7.3059e-02, -7.3547e-02, -7.0435e-02,  ..., -7.0435e-02,
           -7.1655e-02, -6.8726e-02],
          ...,
          [-7.2266e-02, -7.2693e-02, -6.9519e-02,  ..., -6.9580e-02,
           -7.1228e-02, -6.8481e-02],
          [-7.4036e-02, -7.2449e-02, -7.0801e-02,  ..., -6.9946e-02,
           -7.2021e-02, -6.8848e-02],
          [-7.2754e-02, -6.7810e-02, -6.6223e-02,  ..., -6.6528e-02,
           -6.7078e-02, -6.3965e-02]],

         [[-4.8103e-03, -4.7607e-03, -4.0703e-03,  ..., -4.2229e-03,
           -3.9940e-03, -4.3488e-03],
          [-5.9967e-03, -2.1973e-03,  7.7915e-04,  ...,  4.7016e-04,
            7.7295e-04,  8.6784e-05],
          [-6.2561e-03, -3.0174e-03,  6.2656e-04,  ..., -4.8637e-05,
            5.3024e-04,  8.0013e-04],
          ...,
          [-5.8899e-03, -1.5860e-03,  1.9331e-03,  ...,  4.7636e-04,
            9.0313e-04,  1.1959e-03],
          [-4.4518e-03, -7.7724e-04,  2.3994e-03,  ...,  1.0509e-03,
            1.4935e-03,  4.6062e-04],
          [-4.4060e-03, -1.1349e-04,  3.0861e-03,  ...,  2.1420e-03,
            3.0632e-03,  1.6823e-03]]]], device='cuda:0', dtype=torch.float16)

[narendasan]

@peri044 can you take a look?

[AIManifest]

Thank you for your response! Much appreciated friend. Looking forward to any insightful inputs:)

[AIManifest]

Hey friends, kindly following up on any insights? @peri044 @narendasan