minor fix: Small changes

Author: gs-olive

py/torch_tensorrt/dynamo/test/test_dynamo_backend.py

@@ -1,6 +1,7 @@
 import torch
 import timm
 import pytest
+import unittest
 
 import torch_tensorrt as torchtrt
 import torchvision.models as models
@@ -12,6 +13,8 @@
     cosine_similarity,
 )
 
+assertions = unittest.TestCase()
+
 
 @pytest.mark.unit
 def test_resnet18(ir):
@@ -32,9 +35,9 @@ def test_resnet18(ir):
 
     trt_mod = torchtrt.compile(model, **compile_spec)
     cos_sim = cosine_similarity(model(input), trt_mod(input))
-    assert (
+    assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
-        f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env
@@ -63,9 +66,9 @@ def test_mobilenet_v2(ir):
 
     trt_mod = torchtrt.compile(model, **compile_spec)
     cos_sim = cosine_similarity(model(input), trt_mod(input))
-    assert (
+    assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
-        f"Mobilenet v2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        msg=f"Mobilenet v2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env
@@ -94,9 +97,9 @@ def test_efficientnet_b0(ir):
 
     trt_mod = torchtrt.compile(model, **compile_spec)
     cos_sim = cosine_similarity(model(input), trt_mod(input))
-    assert (
+    assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
-        f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        msg=f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env
@@ -138,9 +141,9 @@ def test_bert_base_uncased(ir):
     for key in model_outputs.keys():
         out, trt_out = model_outputs[key], trt_model_outputs[key]
         cos_sim = cosine_similarity(out, trt_out)
-        assert (
+        assertions.assertTrue(
             cos_sim > COSINE_THRESHOLD,
-            f"HF BERT base-uncased TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+            msg=f"HF BERT base-uncased TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
         )
 
     # Clean up model env
@@ -169,9 +172,9 @@ def test_resnet18_half(ir):
 
     trt_mod = torchtrt.compile(model, **compile_spec)
     cos_sim = cosine_similarity(model(input), trt_mod(input))
-    assert (
+    assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
-        f"Resnet18 Half TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        msg=f"Resnet18 Half TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -159,36 +159,32 @@ def acc_ops_conv_transposend(
             input_val.shape[1] != -1
         ), "Channel dim can't be dynamic for transpose convolution."
 
-    if not isinstance(kwargs["bias"], TRTTensor):
-        if kwargs["bias"] is not None and not isinstance(kwargs["bias"], torch.Tensor):
-            raise RuntimeError(
-                f"linear {name} has bias of type {type(kwargs['bias'])}, Expect Optional[Tensor]"
-            )
-        bias = to_numpy(kwargs["bias"])  # type: ignore[arg-type]
-    else:
-        bias = kwargs["bias"]
+    # for now we'll assume bias is constant Tensor or None,
+    # and bias being ITensor is not supported in TensorRT api
+    # right now
+    if kwargs["bias"] is not None and not isinstance(kwargs["bias"], torch.Tensor):
+        raise RuntimeError(
+            f"ConvTranspose {name} has bias of type {type(kwargs['bias'])}, Expect Optional[Tensor]"
+        )
+    bias = to_numpy(kwargs["bias"])  # type: ignore[arg-type]
 
     if network.has_explicit_precision or isinstance(kwargs["weight"], TRTTensor):
         weight = get_trt_tensor(network, kwargs["weight"], f"{name}_weight")
         weight_shape = tuple(kwargs["weight"].shape)  # type: ignore[union-attr]
         # will need to use uninitialized weight and set it later to support
         # ITensor weights
+        dummy_weight = trt.Weights()
 
         # nn.ConvTranspose2d/3d weight size is (in_channels, out_channels/groups, kernel_0, kernel_1, [kernel_2])
         layer = network.add_deconvolution_nd(
             input=input_val,
             num_output_maps=weight.shape[1] * kwargs["groups"],
             kernel_shape=weight.shape[2:],
-            kernel=trt.Weights(),
-            bias=trt.Weights() if isinstance(bias, TRTTensor) else bias,
+            kernel=dummy_weight,
+            bias=bias,
         )
 
         layer.set_input(1, weight)
-
-        # If the bias is a TRTTensor, set it as an input of the layer
-        if isinstance(bias, TRTTensor):
-            bias = get_trt_tensor(network, bias, f"{name}_bias")
-            layer.set_input(2, bias)
     else:
         if not isinstance(kwargs["weight"], torch.Tensor):
             raise RuntimeError(

py/torch_tensorrt/fx/converters/impl/convolution.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Union
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
@@ -62,21 +62,8 @@ def convNd(
         # Transform the bias constant into a Numpy array
         bias = to_numpy(bias)
 
-        # Prepend new dimension (unsqueeze) if the convolution is 1d
-        if is_conv1d:
-            bias = np.expand_dims(bias, 0)
-
     elif isinstance(bias, TRTTensor):
         bias = get_trt_tensor(network, bias, f"{name}_bias")
-        # Prepend new dimension (unsqueeze) if the convolution is 1d
-        if is_conv1d:
-            kwargs = {
-                "input": bias,
-                "dim": 0,
-            }
-            bias = acc_ops_unsqueeze(
-                network, target, tuple(), kwargs, name + "_unsqueeze_bias"
-            )
 
     elif bias is not None:
         raise RuntimeError(