From 7592715944c863a0c1d62fabf1925224efdbe4d4 Mon Sep 17 00:00:00 2001
From: cehongwang <123616592+cehongwang@users.noreply.github.com>
Date: Tue, 2 Jul 2024 11:29:06 -0700
Subject: [PATCH] Implemented basic pipeline for Refitting (#2886)

---
 docsrc/py_api/dynamo.rst              |  1 -
 py/torch_tensorrt/dynamo/_compiler.py |  2 --
 py/torch_tensorrt/dynamo/_refit.py    | 15 +++++++++------
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/docsrc/py_api/dynamo.rst b/docsrc/py_api/dynamo.rst
index 0f3e32f2f7..8c6ff7f4a0 100644
--- a/docsrc/py_api/dynamo.rst
+++ b/docsrc/py_api/dynamo.rst
@@ -26,7 +26,6 @@ Functions
 
 .. autofunction:: refit_module_weights
 
-
 Classes
 --------
 
diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index d562277aeb..7a13c303a5 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -176,8 +176,6 @@ def compile(
         else:
             make_refitable = kwargs["refit"]
 
-    if kwarg_inputs is None:
-        kwarg_inputs = {}
     engine_capability = EngineCapability._from(engine_capability)
 
     if torch_executed_modules is not None and torch_executed_modules:
diff --git a/py/torch_tensorrt/dynamo/_refit.py b/py/torch_tensorrt/dynamo/_refit.py
index 569fc2db55..98cab802c9 100644
--- a/py/torch_tensorrt/dynamo/_refit.py
+++ b/py/torch_tensorrt/dynamo/_refit.py
@@ -6,6 +6,10 @@
 from typing import Any, Sequence, Tuple
 
 import numpy as np
+<<<<<<< HEAD
+=======
+import tensorrt as trt
+>>>>>>> 9f46d3940 (Implemented basic pipeline for Refitting (#2886))
 import torch
 from torch.export import ExportedProgram
 from torch_tensorrt._enums import dtype
@@ -42,8 +46,11 @@
 )
 from torch_tensorrt.logging import TRT_LOGGER
 
+<<<<<<< HEAD
 import tensorrt as trt
 
+=======
+>>>>>>> 9f46d3940 (Implemented basic pipeline for Refitting (#2886))
 logger = logging.getLogger(__name__)
 
 
@@ -96,16 +103,12 @@ def construct_refit_mapping(
         layer_type: str = layer.type.name
         if layer_type in MODULE_MAP:
             # Cast the parent class to child class to access attributes
-            # For example: ILayer does not have ILayer.kernel/ILayer.bias
+            # For example: ILayer does not have ILayer.kernal/ILayer.bias
             # So we cast it to IConvolutionLayer and access the attributes
             layer.__class__ = MODULE_MAP[layer_type][0]
             for weight_type, weight_name in MODULE_MAP[layer_type][1]:
                 weight = layer.__getattribute__(weight_type).copy()
-                weight_dtype_opt = dtype.try_from(weight.dtype)
-                assert (
-                    weight_dtype_opt is not None
-                ), f"Weights {weight_name} has unsupported dtype {weight.dtype}"
-                weight_dtype = weight_dtype_opt.to(trt.DataType)
+                weight_dtype = dtype.try_from(weight.dtype).to(trt.DataType)
                 weight_map[f"{layer.name} {weight_name}"] = (
                     weight,
                     weight_dtype,