Revised according to comments

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/_ConversionContext.py

@@ -29,6 +29,15 @@ class ConversionContext:
     )
 
     def record_weight(self, name: str, weight: torch.Tensor) -> None:
+        """
+        Record the weight and name for refitting and CPU reference.
+        For the refit map, the key is the weight name that appears in the TRT engine and the value is the weight tensor.
+        For the CPU reference holder, we need to hold the reference to the weight tensor until the whole compilation process is complete.
+
+        Args:
+            name: Name of the weight
+            weight: Weight to record
+        """
         self.weight_refit_map[name] = weight
         self.cpu_weights_reference_holder[name + " CPU_REFERENCE"] = weight
 

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -345,7 +345,7 @@ def to_trt_weights(
         - Input tensors are made contiguous before conversion
         - Data type is preserved from the original tensor/array
     """
-    if isinstance(value, np.ndarray):
+    if not isinstance(value, torch.Tensor):
         raise AssertionError(
             f"to_trt_weights can only be called on torch.Tensor, got an object of type: {type(value)}"
         )
@@ -355,10 +355,10 @@ def to_trt_weights(
     supported_weight_types = ["KERNEL", "BIAS", "CONSTANT"]
     assert (
         layer_type_name in supported_layer_types
-    ), f"Unsupported layer type: {layer_type_name}. Please add the layer type to this function to enable refitting."
+    ), f"Encountered unsupported layer type: {layer_type_name}. Supported types are: {supported_layer_types}. Manually calling to_trt_weights with a custom layer type is not intended for general use."
     assert (
         weight_type_name in supported_weight_types
-    ), f"Unsupported weight type: {weight_type_name}. Please add the weight type to this function to enable refitting."
+    ), f"Encountered unsupported weight type: {weight_type_name}. Supported types are: {supported_weight_types}. Manually calling to_trt_weights with a custom weight type is not intended for general use."
 
     if weight_type_name == "CONSTANT" and layer_type_name == "CONSTANT":
         weight_name = f"{name} CONSTANT"

py/torch_tensorrt/dynamo/conversion/impl/conv.py

@@ -55,7 +55,15 @@ def convNd(
     # Process bias terms
     if isinstance(bias, (torch.Tensor, np.ndarray)):
         bias = to_torch(bias, dtype=input.dtype)
-        bias = get_trt_tensor(ctx, bias, f"{name}_bias")
+        bias = to_trt_weights(
+            ctx,
+            bias,
+            name,
+            layer_type_name="CONVOLUTION",
+            weight_type_name="BIAS",
+            target=target,
+            source_ir=source_ir,
+        )
 
     elif isinstance(bias, TRTTensor):
         bias = get_trt_tensor(ctx, bias, f"{name}_bias")

py/torch_tensorrt/dynamo/conversion/impl/deconv.py

@@ -55,7 +55,15 @@ def deconvNd(
     if isinstance(bias, (torch.Tensor, np.ndarray)):
         # Transform the bias constant into a Numpy array
         bias = to_torch(bias, dtype=input.dtype)
-        bias = get_trt_tensor(ctx, bias, f"{name}_bias")
+        bias = to_trt_weights(
+            ctx,
+            bias,
+            name,
+            layer_type_name="CONVOLUTION",
+            weight_type_name="BIAS",
+            target=target,
+            source_ir=source_ir,
+        )
 
     elif isinstance(bias, TRTTensor):
         bias = get_trt_tensor(ctx, bias, f"{name}_bias")