#!/usr/bin/python

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import time
import json
import torch
import argparse
import statistics
from collections import Counter

torch_type_to_triton_type = {
    torch.bool: "TYPE_BOOL",
    torch.int8: "TYPE_INT8",
    torch.int16: "TYPE_INT16",
    torch.int32: "TYPE_INT32",
    torch.int64: "TYPE_INT64",
    torch.uint8: "TYPE_UINT8",
    torch.float16: "TYPE_FP16",
    torch.float32: "TYPE_FP32",
    torch.float64: "TYPE_FP64",
}

CONFIG_TEMPLATE = r"""
name: "{model_name}"
platform: "{platform}"
max_batch_size: {max_batch_size}
input [
    {spec_inputs}
]
output [
    {spec_outputs}
]
{dynamic_batching}
{model_optimizations}
instance_group [
    {{
        count: {engine_count}
        kind: KIND_GPU
        gpus: [ {gpu_list} ]
    }}
]"""

INPUT_TEMPLATE = r"""
{{
    name: "input__{num}"
    data_type: {type}
    dims: {dims}
    {reshape}
}},"""

OUTPUT_TEMPLATE = r""" 
{{
    name: "output__{num}"
    data_type: {type}
    dims: {dims}
    {reshape}
}},"""

MODEL_OPTIMIZATION_TEMPLATE = r"""
optimization {{
  {execution_accelerator}
  cuda {{
    graphs: {capture_cuda_graph}
  }}
}}"""

EXECUTION_ACCELERATOR_TEMPLATE = r"""
  execution_accelerators {{
    gpu_execution_accelerator: [
      {{
        name: "tensorrt"
      }}
    ]
  }},"""


def remove_empty_lines(text):
    """ removes empty lines from text, returns the result """
    ret = "".join([s for s in text.strip().splitlines(True) if s.strip()])
    return ret


def create_deployer(argv):
    """ takes a list of arguments, returns a deployer object and the list of unused arguments """
    parser = argparse.ArgumentParser()
    # required args
    method = parser.add_mutually_exclusive_group(required=True)
    method.add_argument(
        "--ts-script",
        action="store_true",
        help="convert to torchscript using torch.jit.script",
    )
    method.add_argument(
        "--ts-trace",
        action="store_true",
        help="convert to torchscript using torch.jit.trace",
    )
    method.add_argument(
        "--onnx", action="store_true", help="convert to onnx using torch.onnx.export"
    )
    method.add_argument(
        "--trt", action="store_true", help="convert to trt using tensorrt"
    )
    # triton related args
    arguments = parser.add_argument_group("triton related flags")
    arguments.add_argument(
        "--triton-no-cuda", action="store_true", help="Use the CPU for tracing."
    )
    arguments.add_argument(
        "--triton-model-name",
        type=str,
        default="model",
        help="exports to appropriate directory structure for TRITON",
    )
    arguments.add_argument(
        "--triton-model-version",
        type=int,
        default=1,
        help="exports to appropriate directory structure for TRITON",
    )
    arguments.add_argument(
        "--triton-max-batch-size",
        type=int,
        default=8,
        help="Specifies the 'max_batch_size' in the TRITON model config.\
                                  See the TRITON documentation for more info.",
    )
    arguments.add_argument(
        "--triton-dyn-batching-delay",
        type=float,
        default=0,
        help="Determines the dynamic_batching queue delay in milliseconds(ms) for\
                                  the TRITON model config. Use '0' or '-1' to specify static batching.\
                                  See the TRITON documentation for more info.",
    )
    arguments.add_argument(
        "--triton-engine-count",
        type=int,
        default=1,
        help="Specifies the 'instance_group' count value in the TRITON model config.\
                                  See the TRITON documentation for more info.",
    )
    arguments.add_argument(
        "--save-dir", type=str, default="./triton_models", help="Saved model directory"
    )
    # optimization args
    arguments = parser.add_argument_group("optimization flags")
    arguments.add_argument(
        "--max_workspace_size",
        type=int,
        default=512 * 1024 * 1024,
        help="set the size of the workspace for trt export",
    )
    arguments.add_argument(
        "--trt-fp16",
        action="store_true",
        help="trt flag ---- export model in mixed precision mode",
    )
    arguments.add_argument(
        "--capture-cuda-graph",
        type=int,
        default=1,
        help="capture cuda graph for obtaining speedup. possible values: 0, 1. default: 1. ",
    )

    # remainder args
    arguments.add_argument(
        "model_arguments",
        nargs=argparse.REMAINDER,
        help="arguments that will be ignored by deployer lib and will be forwarded to your deployer script",
    )
    #
    args = parser.parse_args(argv)
    deployer = Deployer(args)
    #
    return deployer, args.model_arguments[1:]


class DeployerLibrary:
    def __init__(self, args):
        self.args = args
        self.platform = None

    def set_platform(self, platform):
        """ sets the platform
            :: platform :: "pytorch_libtorch" or "onnxruntime_onnx" or "tensorrt_plan"
        """
        self.platform = platform

    def build_trt_engine(self, model_file, shapes):
        """ takes a path to an onnx file, and shape information, returns a trt engine
            :: model_file :: path to an onnx model
            :: shapes :: dictionary containing min shape, max shape, opt shape for the trt engine
        """
        import tensorrt as trt

        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        builder = trt.Builder(TRT_LOGGER)
        builder.fp16_mode = self.args.trt_fp16
        builder.max_batch_size = self.args.triton_max_batch_size
        #
        config = builder.create_builder_config()
        config.max_workspace_size = self.args.max_workspace_size
        if self.args.trt_fp16:
            config.flags |= 1 << int(trt.BuilderFlag.FP16)
        profile = builder.create_optimization_profile()
        for s in shapes:
            profile.set_shape(s["name"], min=s["min"], opt=s["opt"], max=s["max"])
        config.add_optimization_profile(profile)
        explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        network = builder.create_network(explicit_batch)
        #
        with trt.OnnxParser(network, TRT_LOGGER) as parser:
            with open(model_file, "rb") as model:
                parser.parse(model.read())
                for i in range(parser.num_errors):
                    e = parser.get_error(i)
                    print("||||e", e)
                engine = builder.build_engine(network, config=config)
        return engine

    def load_engine(self, engine_filepath):
        """ loads a trt engine from engine_filepath, returns it """
        import tensorrt as trt

        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        with open(engine_filepath, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
        return engine

    def prepare_inputs(self, dataloader, device):
        """ load sample inputs to device """
        inputs = []
        for batch in dataloader:
            if type(batch) is torch.Tensor:
                batch_d = batch.to(device)
                batch_d = (batch_d,)
                inputs.append(batch_d)
            else:
                batch_d = []
                for x in batch:
                    assert type(x) is torch.Tensor, "input is not a tensor"
                    batch_d.append(x.to(device))
                batch_d = tuple(batch_d)
                inputs.append(batch_d)
        return inputs

    def get_list_of_shapes(self, l, fun):
        """ returns the list of min/max shapes, depending on fun
            :: l :: list of tuples of tensors
            :: fun :: min or max
        """
        tensor_tuple = l[0]
        shapes = [list(x.shape) for x in tensor_tuple]
        for tensor_tuple in l:
            assert len(tensor_tuple) == len(
                shapes
            ), "tensors with varying shape lengths are not supported"
            for i, x in enumerate(tensor_tuple):
                for j in range(len(x.shape)):
                    shapes[i][j] = fun(shapes[i][j], x.shape[j])
        return shapes  # a list of shapes

    def get_tuple_of_min_shapes(self, l):
        """ returns the tuple of min shapes 
            :: l :: list of tuples of tensors """
        shapes = self.get_list_of_shapes(l, min)
        min_batch = 1
        shapes = [[min_batch, *shape[1:]] for shape in shapes]
        shapes = tuple(shapes)
        return shapes  # tuple of min shapes

    def get_tuple_of_max_shapes(self, l):
        """ returns the tuple of max shapes 
            :: l :: list of tuples of tensors """
        shapes = self.get_list_of_shapes(l, max)
        max_batch = max(2, shapes[0][0])
        shapes = [[max_batch, *shape[1:]] for shape in shapes]
        shapes = tuple(shapes)
        return shapes  # tuple of max shapes

    def get_tuple_of_opt_shapes(self, l):
        """ returns the tuple of opt shapes 
            :: l :: list of tuples of tensors """
        counter = Counter()
        for tensor_tuple in l:
            shapes = [tuple(x.shape) for x in tensor_tuple]
            shapes = tuple(shapes)
            counter[shapes] += 1
        shapes = counter.most_common(1)[0][0]
        return shapes  # tuple of most common occuring shapes

    def get_tuple_of_dynamic_shapes(self, l):
        """ returns a tuple of dynamic shapes: variable tensor dimensions 
            (for ex. batch size) occur as -1 in the tuple
            :: l :: list of tuples of tensors """
        tensor_tuple = l[0]
        shapes = [list(x.shape) for x in tensor_tuple]
        for tensor_tuple in l:
            err_msg = "tensors with varying shape lengths are not supported"
            assert len(tensor_tuple) == len(shapes), err_msg
            for i, x in enumerate(tensor_tuple):
                for j in range(len(x.shape)):
                    if shapes[i][j] != x.shape[j] or j == 0:
                        shapes[i][j] = -1
        shapes = tuple(shapes)
        return shapes  # tuple of dynamic shapes

    def run_models(self, models, inputs):
        """ run the models on inputs, return the outputs and execution times """
        ret = []
        for model in models:
            torch.cuda.synchronize()
            time_start = time.time()
            outputs = []
            for input in inputs:
                with torch.no_grad():
                    output = model(*input)
                if type(output) is torch.Tensor:
                    output = [output]
                outputs.append(output)
            torch.cuda.synchronize()
            time_end = time.time()
            t = time_end - time_start
            ret.append(outputs)
            ret.append(t)
        return ret

    def compute_tensor_stats(self, tensor):
        return {
            "std": tensor.std().item(),
            "mean": tensor.mean().item(),
            "max": tensor.max().item(),
            "min": tensor.min().item(),
        }

    def compute_errors(self, outputs_A, outputs_B):
        """ returns dictionary with errors statistics """
        device = outputs_A[0][0][0].device
        dtype = outputs_A[0][0][0].dtype
        x_values = torch.zeros(0, device=device, dtype=dtype)
        y_values = torch.zeros(0, device=device, dtype=dtype)
        d_values = torch.zeros(0, device=device, dtype=dtype)
        for output_A, output_B in zip(outputs_A, outputs_B):
            for x, y in zip(output_A, output_B):
                d = abs(x - y)
                x_values = torch.cat((x_values, x), 0)
                y_values = torch.cat((y_values, y), 0)
                d_values = torch.cat((d_values, d), 0)
        Error_stats = {
            "Original": self.compute_tensor_stats(x_values),
            "Converted": self.compute_tensor_stats(y_values),
            "Absolute difference": self.compute_tensor_stats(d_values),
        }
        return Error_stats

    def print_errors(self, Error_stats):
        """ print various statistcs of Linf errors """
        print()
        print("conversion correctness test results")
        print("-----------------------------------")
        import pandas as pd

        print(pd.DataFrame(Error_stats))

    def write_config(
        self, config_filename, input_shapes, input_types, output_shapes, output_types
    ):
        """ writes TRTIS config file 
            :: config_filename :: the file to write the config file into
            :: input_shapes :: tuple of dynamic shapes of the input tensors
            :: input_types :: tuple of torch types of the input tensors
            :: output_shapes :: tuple of dynamic shapes of the output tensors
            :: output_types :: tuple of torch types of the output tensors
        """
        assert self.platform is not None, "error - platform is not set"

        config_template = CONFIG_TEMPLATE
        input_template = INPUT_TEMPLATE
        optimization_template = MODEL_OPTIMIZATION_TEMPLATE
        accelerator_template = EXECUTION_ACCELERATOR_TEMPLATE

        spec_inputs = r""""""
        for i, (shape, typ) in enumerate(zip(input_shapes, input_types)):
            d = {
                "num": str(i),
                "type": torch_type_to_triton_type[typ],
                "dims": str([1])
                if len(shape) == 1
                else str(list(shape)[1:]),  # first dimension is the batch size
            }
            d["reshape"] = "reshape: { shape: [ ] }" if len(shape) == 1 else ""
            spec_inputs += input_template.format_map(d)
        spec_inputs = spec_inputs[:-1]

        output_template = OUTPUT_TEMPLATE
        spec_outputs = r""""""
        for i, (shape, typ) in enumerate(zip(output_shapes, output_types)):
            d = {
                "num": str(i),
                "type": torch_type_to_triton_type[typ],
                "dims": str([1])
                if len(shape) == 1
                else str(list(shape)[1:]),  # first dimension is the batch size
            }
            d["reshape"] = "reshape: { shape: [ ] }" if len(shape) == 1 else ""
            spec_outputs += output_template.format_map(d)
        spec_outputs = spec_outputs[:-1]

        batching_str = ""
        max_batch_size = self.args.triton_max_batch_size

        if self.args.triton_dyn_batching_delay >= 0:
            # Use only full and half full batches
            pref_batch_size = [int(max_batch_size / 2.0), max_batch_size]

            if self.args.triton_dyn_batching_delay > 0:
                dyn_batch_delay_str = f"max_queue_delay_microseconds: {int(self.args.triton_dyn_batching_delay * 1000.0)}"
            else:
                dyn_batch_delay_str = ""

            batching_str = r"""
dynamic_batching {{
    preferred_batch_size: [{0}]
    {1}
}}""".format(
                ", ".join([str(x) for x in pref_batch_size]), dyn_batch_delay_str
            )

        accelerator_str = ""

        d = {
            "execution_accelerator": accelerator_str,
            "capture_cuda_graph": str(self.args.capture_cuda_graph),
        }
        optimization_str = optimization_template.format_map(d)

        config_values = {
            "model_name": self.args.triton_model_name,
            "platform": self.platform,
            "max_batch_size": max_batch_size,
            "spec_inputs": spec_inputs,
            "spec_outputs": spec_outputs,
            "dynamic_batching": batching_str,
            "model_optimizations": optimization_str,
            "gpu_list": ", ".join([str(x) for x in range(torch.cuda.device_count())]),
            "engine_count": self.args.triton_engine_count,
        }

        # write config
        with open(config_filename, "w") as file:
            final_config_str = config_template.format_map(config_values)
            final_config_str = remove_empty_lines(final_config_str)
            file.write(final_config_str)


class Deployer:
    def __init__(self, args):
        self.args = args
        self.lib = DeployerLibrary(args)

    def deploy(self, dataloader, model):
        """ deploy the model and test for correctness with dataloader """
        if self.args.ts_script or self.args.ts_trace:
            self.lib.set_platform("pytorch_libtorch")
            print(
                "deploying model "
                + self.args.triton_model_name
                + " in format "
                + self.lib.platform
            )
            self.to_triton_torchscript(dataloader, model)
        elif self.args.onnx:
            self.lib.set_platform("onnxruntime_onnx")
            print(
                "deploying model "
                + self.args.triton_model_name
                + " in format "
                + self.lib.platform
            )
            self.to_triton_onnx(dataloader, model)
        elif self.args.trt:
            self.lib.set_platform("tensorrt_plan")
            print(
                "deploying model "
                + self.args.triton_model_name
                + " in format "
                + self.lib.platform
            )
            self.to_triton_trt(dataloader, model)
        else:
            assert False, "error"
        print("done")

    def to_triton_trt(self, dataloader, model):
        """ export the model to trt and test correctness on dataloader """
        import tensorrt as trt

        # setup device
        if self.args.triton_no_cuda:
            device = torch.device("cpu")
        else:
            device = torch.device("cuda")

        # prepare model
        model.to(device)
        model.eval()
        assert not model.training, "internal error - model should be in eval() mode! "

        # prepare inputs
        inputs = self.lib.prepare_inputs(dataloader, device)

        # generate outputs
        outputs = []
        for input in inputs:
            with torch.no_grad():
                output = model(*input)
            if type(output) is torch.Tensor:
                output = [output]
            outputs.append(output)

        # generate input shapes - dynamic tensor shape support
        input_shapes = self.lib.get_tuple_of_dynamic_shapes(inputs)

        # generate output shapes - dynamic tensor shape support
        output_shapes = self.lib.get_tuple_of_dynamic_shapes(outputs)

        # generate input types
        input_types = [x.dtype for x in inputs[0]]

        # generate output types
        output_types = [x.dtype for x in outputs[0]]

        # get input names
        rng = range(len(input_types))
        input_names = ["input__" + str(num) for num in rng]

        # get output names
        rng = range(len(output_types))
        output_names = ["output__" + str(num) for num in rng]

        # prepare save path
        model_folder = os.path.join(self.args.save_dir, self.args.triton_model_name)
        version_folder = os.path.join(model_folder, str(self.args.triton_model_version))
        if not os.path.exists(version_folder):
            os.makedirs(version_folder)
        final_model_path = os.path.join(version_folder, "model.plan")

        # get indices of dynamic input and output shapes
        dynamic_axes = {}
        for input_name, shape in zip(input_names, input_shapes):
            dynamic_axes[input_name] = [i for i, x in enumerate(shape) if x == -1]
        for output_name, shape in zip(output_names, output_shapes):
            dynamic_axes[output_name] = [i for i, x in enumerate(shape) if x == -1]

        # export the model to onnx first
        with torch.no_grad():
            torch.onnx.export(
                model,
                inputs[0],
                final_model_path,
                verbose=False,
                input_names=input_names,
                output_names=output_names,
                dynamic_axes=dynamic_axes,
                opset_version=11,
            )

        # get shapes
        min_shapes = self.lib.get_tuple_of_min_shapes(inputs)
        opt_shapes = self.lib.get_tuple_of_opt_shapes(inputs)
        max_shapes = self.lib.get_tuple_of_max_shapes(inputs)

        zipped = zip(input_names, min_shapes, opt_shapes, max_shapes)
        shapes = []
        for name, min_shape, opt_shape, max_shape in zipped:
            d = {"name": name, "min": min_shape, "opt": opt_shape, "max": max_shape}
            shapes.append(d)

        # build trt engine
        engine = self.lib.build_trt_engine(final_model_path, shapes)
        assert engine is not None, " trt export failure "

        # write trt engine
        with open(final_model_path, "wb") as f:
            f.write(engine.serialize())

        # load the model
        engine = self.lib.load_engine(final_model_path)

        class TRT_model:
            def __init__(self, engine, input_names, output_names, output_types, device):
                self.engine = engine
                self.context = self.engine.create_execution_context()
                self.input_names = input_names
                self.output_names = output_names
                self.output_types = output_types
                self.device = device

            def is_dimension_dynamic(self, dim):
                return dim is None or dim <= 0

            def is_shape_dynamic(self, shape):
                return any([self.is_dimension_dynamic(dim) for dim in shape])

            def __call__(self, *inputs):
                # get input shapes
                input_shapes = [x.shape for x in inputs]
                # bindings
                bindings = [None] * self.engine.num_bindings
                # set input shapes, bind input tensors
                zipped = zip(self.input_names, inputs)
                for key, input in zipped:
                    idx = self.engine.get_binding_index(key)
                    bindings[idx] = input.data_ptr()
                    if self.engine.is_shape_binding(idx) and self.is_shape_dynamic(
                        self.context.get_shape(idx)
                    ):
                        self.context.set_shape_input(idx, input)
                    elif self.is_shape_dynamic(self.engine.get_binding_shape(idx)):
                        self.context.set_binding_shape(idx, input.shape)
                assert self.context.all_binding_shapes_specified, "trt error"
                assert self.context.all_shape_inputs_specified, "trt error"
                # calculate output shapes, allocate output tensors and bind them
                outputs = []
                zipped = zip(self.output_names, self.output_types)
                for key, dtype in zipped:
                    idx = self.engine.get_binding_index(key)
                    shape = self.context.get_binding_shape(idx)
                    shape = tuple(shape)
                    assert -1 not in shape, "trt error"
                    tensor = torch.zeros(shape, dtype=dtype, device=self.device)
                    outputs.append(tensor)
                    bindings[idx] = outputs[-1].data_ptr()
                # run inference
                self.context.execute_v2(bindings=bindings)
                # return the result
                if len(outputs) == 1:
                    outputs = outputs[0]
                return outputs

        model_trt = TRT_model(engine, input_names, output_names, output_types, device)

        # run both models on inputs
        assert not model.training, "internal error - model should be in eval() mode! "
        models = (model, model_trt)
        outputs, time_model, outputs_trt, time_model_trt = self.lib.run_models(
            models, inputs
        )

        # check for errors
        Error_stats = self.lib.compute_errors(outputs, outputs_trt)
        self.lib.print_errors(Error_stats)
        print("time of error check of native model: ", time_model, "seconds")
        print("time of error check of trt model: ", time_model_trt, "seconds")
        print()

        # write TRTIS config
        config_filename = os.path.join(model_folder, "config.pbtxt")
        self.lib.write_config(
            config_filename, input_shapes, input_types, output_shapes, output_types
        )

    def name_onnx_nodes(self, model_path):
        """
        Name all unnamed nodes in ONNX model
            parameter model_path: path  ONNX model
            return: none
        """
        model = onnx.load(model_path)
        node_id = 0
        for node in model.graph.node:
            if len(node.name) == 0:
                node.name = "unnamed_node_%d" % node_id
            node_id += 1
        # This check partially validates model
        onnx.checker.check_model(model)
        onnx.save(model, model_path)
        # Only inference really checks ONNX model for some issues
        # like duplicated node names
        onnxruntime.InferenceSession(model_path, None)

    def to_triton_onnx(self, dataloader, model):
        """ export the model to onnx and test correctness on dataloader """
        import onnx as local_onnx

        global onnx
        onnx = local_onnx
        import onnxruntime as local_onnxruntime

        global onnxruntime
        onnxruntime = local_onnxruntime
        # setup device
        if self.args.triton_no_cuda:
            device = torch.device("cpu")
        else:
            device = torch.device("cuda")

        # prepare model
        model.to(device)
        model.eval()
        assert not model.training, "internal error - model should be in eval() mode! "

        # prepare inputs
        inputs = self.lib.prepare_inputs(dataloader, device)

        # generate outputs
        outputs = []
        for input in inputs:
            with torch.no_grad():
                output = model(*input)
            if type(output) is torch.Tensor:
                output = [output]
            outputs.append(output)

        # generate input shapes - dynamic tensor shape support
        input_shapes = self.lib.get_tuple_of_dynamic_shapes(inputs)

        # generate output shapes - dynamic tensor shape support
        output_shapes = self.lib.get_tuple_of_dynamic_shapes(outputs)

        # generate input types
        input_types = [x.dtype for x in inputs[0]]

        # generate output types
        output_types = [x.dtype for x in outputs[0]]

        # get input names
        rng = range(len(input_types))
        input_names = ["input__" + str(num) for num in rng]

        # get output names
        rng = range(len(output_types))
        output_names = ["output__" + str(num) for num in rng]

        # prepare save path
        model_folder = os.path.join(self.args.save_dir, self.args.triton_model_name)
        version_folder = os.path.join(model_folder, str(self.args.triton_model_version))
        if not os.path.exists(version_folder):
            os.makedirs(version_folder)
        final_model_path = os.path.join(version_folder, "model.onnx")

        # get indices of dynamic input and output shapes
        dynamic_axes = {}
        for input_name, input_shape in zip(input_names, input_shapes):
            dynamic_axes[input_name] = [i for i, x in enumerate(input_shape) if x == -1]
        for output_name, output_shape in zip(output_names, output_shapes):
            dynamic_axes[output_name] = [
                i for i, x in enumerate(output_shape) if x == -1
            ]

        # export the model
        assert not model.training, "internal error - model should be in eval() mode! "
        with torch.no_grad():
            torch.onnx.export(
                model,
                inputs[0],
                final_model_path,
                verbose=True,
                input_names=input_names,
                output_names=output_names,
                dynamic_axes=dynamic_axes,
                opset_version=11,
            )

        # syntactic error check
        converted_model = onnx.load(final_model_path)
        # check that the IR is well formed
        onnx.checker.check_model(converted_model)

        # Name unnamed nodes - it helps for some other processing tools
        self.name_onnx_nodes(final_model_path)
        converted_model = onnx.load(final_model_path)

        # load the model
        session = onnxruntime.InferenceSession(final_model_path, None)

        class ONNX_model:
            def __init__(self, session, input_names, device):
                self.session = session
                self.input_names = input_names

            def to_numpy(self, tensor):
                return (
                    tensor.detach().cpu().numpy()
                    if tensor.requires_grad
                    else tensor.cpu().numpy()
                )

            def __call__(self, *inputs):
                inp = [
                    (input_name, inputs[i])
                    for i, input_name in enumerate(self.input_names)
                ]
                inp = {input_name: self.to_numpy(x) for input_name, x in inp}
                outputs = self.session.run(None, inp)
                outputs = [torch.from_numpy(output) for output in outputs]
                outputs = [output.to(device) for output in outputs]
                if len(outputs) == 1:
                    outputs = outputs[0]
                return outputs

        # switch to eval mode
        model_onnx = ONNX_model(session, input_names, device)

        # run both models on inputs
        assert not model.training, "internal error - model should be in eval() mode! "
        models = (model, model_onnx)
        outputs, time_model, outputs_onnx, time_model_onnx = self.lib.run_models(
            models, inputs
        )

        # check for errors
        Error_stats = self.lib.compute_errors(outputs, outputs_onnx)
        self.lib.print_errors(Error_stats)
        print("time of error check of native model: ", time_model, "seconds")
        print("time of error check of onnx model: ", time_model_onnx, "seconds")
        print()

        # write TRTIS config
        config_filename = os.path.join(model_folder, "config.pbtxt")
        self.lib.write_config(
            config_filename, input_shapes, input_types, output_shapes, output_types
        )

    def to_triton_torchscript(self, dataloader, model):
        """ export the model to torchscript and test correctness on dataloader """
        # setup device
        if self.args.triton_no_cuda:
            device = torch.device("cpu")
        else:
            device = torch.device("cuda")

        # prepare model
        model.to(device)
        model.eval()
        assert not model.training, "internal error - model should be in eval() mode! "

        # prepare inputs
        inputs = self.lib.prepare_inputs(dataloader, device)

        # generate input shapes - dynamic tensor shape support
        input_shapes = self.lib.get_tuple_of_dynamic_shapes(inputs)

        # generate input types
        input_types = [x.dtype for x in inputs[0]]

        # prepare save path
        model_folder = os.path.join(self.args.save_dir, self.args.triton_model_name)
        version_folder = os.path.join(model_folder, str(self.args.triton_model_version))
        if not os.path.exists(version_folder):
            os.makedirs(version_folder)
        final_model_path = os.path.join(version_folder, "model.pt")

        # convert the model
        with torch.no_grad():
            if self.args.ts_trace:  # trace it
                model_ts = torch.jit.trace(model, inputs[0])
            if self.args.ts_script:  # script it
                model_ts = torch.jit.script(model)

        # save the model
        torch.jit.save(model_ts, final_model_path)

        # load the model
        model_ts = torch.jit.load(final_model_path)
        model_ts.eval()  # WAR for bug : by default, model_ts gets loaded in training mode

        # run both models on inputs
        assert not model.training, "internal error - model should be in eval() mode! "
        assert (
            not model_ts.training
        ), "internal error - converted model should be in eval() mode! "
        models = (model, model_ts)
        outputs, time_model, outputs_ts, time_model_ts = self.lib.run_models(
            models, inputs
        )

        # check for errors
        Error_stats = self.lib.compute_errors(outputs, outputs_ts)
        self.lib.print_errors(Error_stats)
        print("time of error check of native model: ", time_model, "seconds")
        print("time of error check of ts model: ", time_model_ts, "seconds")
        print()

        # generate output shapes - dynamic tensor shape support
        output_shapes = self.lib.get_tuple_of_dynamic_shapes(outputs)

        # generate output types
        output_types = [x.dtype for x in outputs[0]]

        # now we build the config for TRTIS
        config_filename = os.path.join(model_folder, "config.pbtxt")
        self.lib.write_config(
            config_filename, input_shapes, input_types, output_shapes, output_types
        )