Repository: pytorch/TensorRT (language: python)

Issue:
🐛 [Bug] Importing `torchao` first breaks `torch_tensorrt.dynamo.compile` during `run_decompositions`
##  Bug Description

Importing `torchao` before importing `torch_tensorrt` causes `F.interpolate` to fail during `run_decompositions` with:

`AssertionError: Expected aten.upsample_nearest2d.default to have CompositeImplicitAutograd kernel`

## To Reproduce

```py
import torchao
import torch_tensorrt
import torch
import torch.nn.functional as F

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        _, _, h, w = x.shape
        z = F.interpolate(x, [h*2, w*2])
        return z

model = Model().cuda()
inputs = (torch.randn((2, 4, 8, 8)).cuda(),)

with torch.no_grad():
    ep = torch.export.export(
        model,
        args=inputs,
        strict=True
    )
    with torch_tensorrt.logging.debug():
        trt_gm = torch_tensorrt.dynamo.compile(
            ep,
            inputs,
            reuse_cached_engines=False,
            cache_built_engines=False,
            require_full_compilation=True,
            min_block_size=1,
        )  
```

**Logs**:

```txt
WARNING:torch_tensorrt.dynamo.conversion.aten_ops_converters:Unable to import quantization op. Please install modelopt library (https://github.com/NVIDIA/TensorRT-Model-Optimizer?tab=readme-ov-file#installation) to add support for compiling quantized models
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/projects/scripts/mre/torchao_tensorrt_import.py", line 25, in <module>
    trt_gm = torch_tensorrt.dynamo.compile(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch_tensorrt/dynamo/_compiler.py", line 228, in compile
    exported_program = exported_program.run_decompositions(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/export/exported_program.py", line 116, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/export/exported_program.py", line 1111, in run_decompositions
    return _decompose_exported_program(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/export/exported_program.py", line 654, in _decompose_exported_program
    gm, new_graph_signature = _decompose_and_get_gm_with_new_signature_constants(
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/export/exported_program.py", line 446, in _decompose_and_get_gm_with_new_signature_constants
    gm, graph_signature = aot_export_module(
                          ^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1262, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
                                        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1497, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 524, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 625, in _create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 194, in inner
    flat_f_outs = f(*flat_f_args)
                  ^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 859, in functional_call
    out = PropagateUnbackedSymInts(mod).run(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
                     ^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6495, in run_node
    result = super().run_node(n)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/fx/interpreter.py", line 275, in call_function
    return target(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_ops.py", line 723, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_higher_order_ops/utils.py", line 64, in inner
    return autograd_not_implemented_inner(op, deferred_error, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_higher_order_ops/utils.py", line 37, in autograd_not_implemented_inner
    result = operator(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_ops.py", line 723, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_subclasses/functional_tensor.py", line 449, in __torch_dispatch__
    r = func.decompose(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_ops.py", line 764, in decompose
    return self.py_kernels[dk](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch_tensorrt/dynamo/conversion/aten_ops_converters.py", line 3184, in upsample_nearest2d_vec
    return torch.ops.aten.upsample_nearest2d.default(input, osize)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_ops.py", line 723, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_subclasses/functional_tensor.py", line 449, in __torch_dispatch__
    r = func.decompose(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_ops.py", line 764, in decompose
    return self.py_kernels[dk](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dgcnz/.conda/envs/edge/lib/python3.12/site-packages/torch/_decomp/__init__.py", line 376, in _special_op_to_decompose_cia
    raise AssertionError(
AssertionError: Expected aten.upsample_nearest2d.default to have CompositeImplicitAutograd kernel

While executing %upsample_nearest2d : [num_users=1] = call_function[target=torch.ops.aten.upsample_nearest2d.vec](args = (%x, [16, 16], None), kwargs = {})
Original traceback:
  File "/projects/scripts/mre/torchao_tensorrt_import.py", line 12, in forward
    z = F.interpolate(x, [h*2, w*2])
```


## Expected behavior

The import order between`torch_tensorrt` and `torchao` should not matter.

## Environment

> Build information about Torch-TensorRT can be found by turning on debug messages

 - Torch-TensorRT Version (e.g. 1.0.0): 2.6.0.dev20241008+cu124
 - PyTorch Version (e.g. 1.0): 2.6.0.dev20241009+cu124
 - CPU Architecture: x86_64
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): pip
 - Build command you used (if compiling from source): -
 - Are you using local sources or building from archives: -
 - Python version: 3.12.7
 - CUDA version: 12.4
 - GPU models and configuration: NVIDIA RTX 4060 TI 
 - Any other relevant information:
 - torchao==0.6.0.dev20241009+cu124

## Additional context

If you import `torch_tensorrt` first and then `torchao` the error disappears.


Discussion / hints:
From what I can take from the back trace seems like its a torchao/dyanmo issue but @peri044 and @lanluo-nvidia can you look at this as well as part of the quantization work? 
@dgcnz 
I have tried  the example with the following setup:
```
(torch_tensorrt_release_py39) lanl@d5c2237-lcedt:~/git/py39/TensorRT$ python -m pip list | grep torch
pytorch-triton           3.1.0
torch                    2.5.0+cu124
torch_tensorrt           2.5.0+cu124
torchao                  0.6.0+cu124
torchprofile             0.0.4
torchvision              0.20.0+cu124
```
It is working no matter the import sequence. 
however it does fail on the latest main, which I believe is related to the decomposition changes introduced in torch 2.6.0: https://github.com/pytorch/pytorch/pull/135080
I will take a further look and get the corresponding changes in our torch_tensorrt.

I think it's related to the changes I did at https://github.com/pytorch/TensorRT/pull/2790/files#diff-9cbf43cfb62cdf682eef77f6fd5cabc488bb5a91ff09bfee5d593827542b2476R3044 to not let PyTorch decompose those ops. I may find a way to fix it in one or two days.

Fix patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index f88a1a5e3c..07c8c03697 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -3110,232 +3110,21 @@ def aten_ops_pad(
     )
 
 
-for op in (
-    torch.ops.aten.upsample_nearest1d,
-    torch.ops.aten.upsample_nearest2d,
-    torch.ops.aten.upsample_nearest3d,
-    torch.ops.aten.upsample_linear1d,
-    torch.ops.aten.upsample_bilinear2d,
-    torch.ops.aten.upsample_trilinear3d,
-    torch.ops.aten.upsample_bicubic2d,
-):
-    for key in (
-        torch._C.DispatchKey.Autograd,
-        torch._C.DispatchKey.CompositeImplicitAutograd,
-    ):
-        if key in op.default.py_kernels:
-            del op.default.py_kernels[key]
-        if key in op.vec.py_kernels:
-            del op.vec.py_kernels[key]
-
-
-def upsample_compute_output_size(
-    input_size: torch.Size,
-    output_size: Optional[Sequence[int]],
-    scale_factors: Optional[Sequence[float]],
-) -> Optional[Sequence[int]]:
-    spatial_dimensions = len(input_size) - 2
-
-    if output_size is None and scale_factors is None:
-        raise AssertionError(
-            "Must specify exactly one of output_size and scale_factors"
-        )
-
-    if output_size is not None:
-        torch._check(
-            scale_factors is None,
-            lambda: "Must specify exactly one of output_size and scale_factors",
-        )
-        torch._check(len(output_size) == spatial_dimensions)
-
-    if scale_factors is not None:
-        torch._check(
-            output_size is None,
-            lambda: "Must specify exactly one of output_size and scale_factors",
-        )
-        torch._check(len(scale_factors) == spatial_dimensions)
-        output_size = []
-        for i, s in enumerate(scale_factors):
-            output_size.append(int(input_size[i + 2] * s))
-
-    return output_size
-
-
-@torch.ops.aten.upsample_nearest1d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_nearest1d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_nearest1d.default(input, osize, *scale_factors)
-    return torch.ops.aten.upsample_nearest1d.default(input, osize)
-
-
-@torch.ops.aten.upsample_nearest2d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_nearest2d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_nearest2d.default(input, osize, *scale_factors)
-    return torch.ops.aten.upsample_nearest2d.default(input, osize)
-
-
-@torch.ops.aten.upsample_nearest3d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_nearest3d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_nearest3d.default(input, osize, *scale_factors)
-    return torch.ops.aten.upsample_nearest3d.default(input, osize)
-
-
-@torch.ops.aten.upsample_linear1d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_linear1d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    align_corners: bool,
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_linear1d.default(
-            input, osize, align_corners, *scale_factors
-        )
-    return torch.ops.aten.upsample_linear1d.default(input, osize, align_corners)
-
-
-@torch.ops.aten.upsample_bilinear2d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_bilinear2d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    align_corners: bool,
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_bilinear2d.default(
-            input, osize, align_corners, *scale_factors
-        )
-    return torch.ops.aten.upsample_bilinear2d.default(input, osize, align_corners)
-
-
-@torch.ops.aten.upsample_trilinear3d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_trilinear3d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    align_corners: bool,
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_trilinear3d.default(
-            input, osize, align_corners, *scale_factors
-        )
-    return torch.ops.aten.upsample_trilinear3d.default(input, osize, align_corners)
-
-
-@torch.ops.aten.upsample_bicubic2d.vec.py_impl(
-    torch._C.DispatchKey.CompositeImplicitAutograd
-)
-def upsample_bicubic2d_vec(
-    input: torch.Tensor,
-    output_size: Optional[Sequence[int]],
-    align_corners: bool,
-    scale_factors: Optional[Sequence[float]],
-) -> torch.Tensor:
-    osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
-    if scale_factors is not None:
-        return torch.ops.aten.upsample_bicubic2d.default(
-            input, osize, align_corners, *scale_factors
-        )
-    return torch.ops.aten.upsample_bicubic2d.default(input, osize, align_corners)
-
-
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_nearest1d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_nearest1d.vec, supports_dynamic_shapes=True
 )
-@enforce_tensor_types(
-    {
-        0: (TRTTensor,),
-    }
-)
-def aten_ops_upsample_nearest1d(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.upsample.upsample(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 3 else [args[2]],
-        mode="nearest",
-        align_corners=False,
-    )
-
-
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_nearest2d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_nearest2d.vec, supports_dynamic_shapes=True
 )
-@enforce_tensor_types(
-    {
-        0: (TRTTensor,),
-    }
-)
-def aten_ops_upsample_nearest2d(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.upsample.upsample(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 4 else [args[2], args[3]],
-        mode="nearest",
-        align_corners=False,
-    )
-
-
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_nearest3d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_nearest3d.vec, supports_dynamic_shapes=True
 )
 @enforce_tensor_types(
     {
         0: (TRTTensor,),
     }
 )
-def aten_ops_upsample_nearest3d(
+def aten_ops_upsample_nearest(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
@@ -3348,78 +3137,28 @@ def aten_ops_upsample_nearest3d(
         SourceIR.ATEN,
         name,
         args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 5 else [args[2], args[3], args[4]],
+        size=args_bounds_check(args, 1),
+        scale_factor=args_bounds_check(args, 2),
         mode="nearest",
         align_corners=False,
     )
 
 
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_linear1d.default, supports_dynamic_shapes=True
-)
-@enforce_tensor_types(
-    {
-        0: (TRTTensor,),
-    }
+    torch.ops.aten.upsample_linear1d.vec, supports_dynamic_shapes=True
 )
-def aten_ops_upsample_linear1d(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.upsample.upsample(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 4 else [args[3]],
-        mode="linear",
-        align_corners=args[2],
-    )
-
-
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_bilinear2d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_bilinear2d.vec, supports_dynamic_shapes=True
 )
-@enforce_tensor_types(
-    {
-        0: (TRTTensor,),
-    }
-)
-def aten_ops_upsample_bilinear2d(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.upsample.upsample(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 5 else [args[3], args[4]],
-        mode="bilinear",
-        align_corners=args[2],
-    )
-
-
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_trilinear3d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_trilinear3d.vec, supports_dynamic_shapes=True
 )
 @enforce_tensor_types(
     {
         0: (TRTTensor,),
     }
 )
-def aten_ops_upsample_trilinear3d(
+def aten_ops_upsample_linear(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
@@ -3432,15 +3171,15 @@ def aten_ops_upsample_trilinear3d(
         SourceIR.ATEN,
         name,
         args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 6 else [args[3], args[4], args[5]],
-        mode="trilinear",
+        size=args_bounds_check(args, 1),
+        scale_factor=args_bounds_check(args, 3),
+        mode="linear",
         align_corners=args[2],
     )
 
 
 @dynamo_tensorrt_converter(
-    torch.ops.aten.upsample_bicubic2d.default, supports_dynamic_shapes=True
+    torch.ops.aten.upsample_bicubic2d.vec, supports_dynamic_shapes=True
 )
 @enforce_tensor_types(
     {
@@ -3460,8 +3199,8 @@ def aten_ops_upsample_bicubic2d(
         SourceIR.ATEN,
         name,
         args[0],
-        size=args[1],
-        scale_factor=None if len(args) < 5 else [args[3], args[4]],
+        size=args_bounds_check(args, 1),
+        scale_factor=args_bounds_check(args, 3),
         mode="bicubic",
         align_corners=args[2],
     )
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/upsample.py b/py/torch_tensorrt/dynamo/conversion/impl/upsample.py
index 5d1e281699..247179455c 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/upsample.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/upsample.py
@@ -18,14 +18,14 @@ def upsample(
     source_ir: Optional[SourceIR],
     name: str,
     input: TRTTensor,
-    size: Sequence[int],
+    size: Optional[Sequence[int]],
     scale_factor: Optional[Sequence[float]],
     mode: str,
     align_corners: bool,
 ) -> TRTTensor:
     layer = ctx.net.add_resize(input)
 
-    if scale_factor is not None and all(s is not None for s in scale_factor):
+    if scale_factor is not None:
         layer.scales = [1.0, 1.0] + list(scale_factor)
     else:
         shape = list(input.shape)[:2] + list(size)
diff --git a/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py b/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
index c472c31a84..825be75076 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
@@ -164,6 +164,13 @@
 }
 torch_disabled_decompositions: Set[Union[OpOverload, OpOverloadPacket]] = {
     aten._softmax.default,
+    aten.upsample_nearest1d.vec,
+    aten.upsample_nearest2d.vec,
+    aten.upsample_nearest3d.vec,
+    aten.upsample_linear1d.vec,
+    aten.upsample_bilinear2d.vec,
+    aten.upsample_trilinear3d.vec,
+    aten.upsample_bicubic2d.vec,
 }
 
 
diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
index d195ad81b8..dda014890d 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
@@ -3,7 +3,8 @@
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
-from torch._decomp import _decomp_table_to_post_autograd_aten, register_decomposition
+from torch._decomp import register_decomposition
+from torch._export.utils import _decomp_table_to_post_autograd_aten
 from torch._ops import OpOverload
 from torch_tensorrt.dynamo._defaults import default_device
 from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim
@@ -411,8 +412,15 @@ def get_decompositions(
         return {**CORE_ATEN_DECOMPOSITIONS_FILTERED, **TORCH_TRT_DECOMPOSITIONS}
     else:
         # changes made here due to torch2.6 changes https://github.com/pytorch/pytorch/pull/135080
+        decomp_table = _decomp_table_to_post_autograd_aten()
+        DECOMP_TABLE_FILTERED: Dict[OpOverload, Callable[[Any], Any]] = {
+            decomp: decomp_table[decomp]
+            for decomp in decomp_table
+            if decomp not in torch_disabled_decompositions
+        }
+
         return {
             **ENABLED_TORCH_DECOMPOSITIONS,
-            **_decomp_table_to_post_autograd_aten(),
+            **DECOMP_TABLE_FILTERED,
             **TORCH_TRT_DECOMPOSITIONS,
         }


Test patch:
diff --git a/tests/py/dynamo/conversion/test_upsample_aten.py b/tests/py/dynamo/conversion/test_upsample_aten.py
index 0e33838349..44c4af2a92 100644
--- a/tests/py/dynamo/conversion/test_upsample_aten.py
+++ b/tests/py/dynamo/conversion/test_upsample_aten.py
@@ -10,13 +10,15 @@ class TestUpsampleConverter(DispatchTestCase):
     @parameterized.expand(
         [
             ([7], [3], None),
-            ([7], [10], 1.5),
+            ([7], None, [1.5]),
         ]
     )
-    def test_nearest1d(self, input_size, output_size, scales):
+    def test_nearest1d(self, input_size, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest1d.default(x, output_size, scales)
+                return torch.ops.aten.upsample_nearest1d.vec(
+                    x, output_size, scale_factors
+                )
 
         inputs = [torch.randn([1, 1] + input_size)]
         self.run_test(TestModule(), inputs)
@@ -24,13 +26,15 @@ def forward(self, x):
     @parameterized.expand(
         [
             ([3], None),
-            ([13], 1.5),
+            (None, [1.5]),
         ]
     )
-    def test_nearest1d_dynamic_shape(self, output_size, scales):
+    def test_nearest1d_dynamic_shape(self, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest1d.default(x, output_size, scales)
+                return torch.ops.aten.upsample_nearest1d.vec(
+                    x, output_size, scale_factors
+                )
 
         input_specs = [
             Input(
@@ -44,15 +48,15 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([7, 7], [3, 3], None, None),
-            ([7, 7], [3, 10], 0.5, 1.5),
+            ([7, 7], [3, 3], None),
+            ([7, 7], None, [0.5, 1.5]),
         ]
     )
-    def test_nearest2d(self, input_size, output_size, scales_h, scales_w):
+    def test_nearest2d(self, input_size, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest2d.default(
-                    x, output_size, scales_h, scales_w
+                return torch.ops.aten.upsample_nearest2d.vec(
+                    x, output_size, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -60,15 +64,15 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([3, 3], None, None),
-            ([4, 13], 0.5, 1.5),
+            ([3, 3], None),
+            (None, [0.5, 1.5]),
         ]
     )
-    def test_nearest2d_dynamic_shape(self, output_size, scales_h, scales_w):
+    def test_nearest2d_dynamic_shape(self, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest2d.default(
-                    x, output_size, scales_h, scales_w
+                return torch.ops.aten.upsample_nearest2d.vec(
+                    x, output_size, scale_factors
                 )
 
         input_specs = [
@@ -83,15 +87,15 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([7, 7, 7], [3, 3, 3], None, None, None),
-            ([7, 7, 7], [3, 7, 10], 0.5, 1.0, 1.5),
+            ([7, 7, 7], [3, 3, 3], None),
+            ([7, 7, 7], None, [0.5, 1.0, 1.5]),
         ]
     )
-    def test_nearest3d(self, input_size, output_size, scales_d, scales_h, scales_w):
+    def test_nearest3d(self, input_size, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest3d.default(
-                    x, output_size, scales_d, scales_h, scales_w
+                return torch.ops.aten.upsample_nearest3d.vec(
+                    x, output_size, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -99,15 +103,15 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([3, 3, 3], None, None, None),
-            ([4, 9, 13], 0.5, 1.0, 1.5),
+            ([3, 3, 3], None),
+            (None, [0.5, 1.0, 1.5]),
         ]
     )
-    def test_nearest3d_dynamic_shape(self, output_size, scales_d, scales_h, scales_w):
+    def test_nearest3d_dynamic_shape(self, output_size, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_nearest3d.default(
-                    x, output_size, scales_d, scales_h, scales_w
+                return torch.ops.aten.upsample_nearest3d.vec(
+                    x, output_size, scale_factors
                 )
 
         input_specs = [
@@ -124,15 +128,15 @@ def forward(self, x):
         [
             ([7], [3], True, None),
             ([7], [3], False, None),
-            ([7], [10], True, 1.5),
-            ([7], [10], False, 1.5),
+            ([7], None, True, [1.5]),
+            ([7], None, False, [1.5]),
         ]
     )
-    def test_linear1d(self, input_size, output_size, align_corners, scales):
+    def test_linear1d(self, input_size, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_linear1d.default(
-                    x, output_size, align_corners, scales
+                return torch.ops.aten.upsample_linear1d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -142,15 +146,15 @@ def forward(self, x):
         [
             ([3], True, None),
             ([3], False, None),
-            ([13], True, 1.5),
-            ([13], False, 1.5),
+            (None, True, [1.5]),
+            (None, False, [1.5]),
         ]
     )
-    def test_linear1d_dynamic_shape(self, output_size, align_corners, scales):
+    def test_linear1d_dynamic_shape(self, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_linear1d.default(
-                    x, output_size, align_corners, scales
+                return torch.ops.aten.upsample_linear1d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         input_specs = [
@@ -165,19 +169,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([7, 7], [3, 3], True, None, None),
-            ([7, 7], [3, 3], False, None, None),
-            ([7, 7], [3, 10], True, 0.5, 1.5),
-            ([7, 7], [3, 10], False, 0.5, 1.5),
+            ([7, 7], [3, 3], True, None),
+            ([7, 7], [3, 3], False, None),
+            ([7, 7], None, True, [0.5, 1.5]),
+            ([7, 7], None, False, [0.5, 1.5]),
         ]
     )
-    def test_bilinear2d(
-        self, input_size, output_size, align_corners, scales_h, scales_w
-    ):
+    def test_bilinear2d(self, input_size, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_bilinear2d.default(
-                    x, output_size, align_corners, scales_h, scales_w
+                return torch.ops.aten.upsample_bilinear2d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -185,19 +187,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([3, 3], True, None, None),
-            ([3, 3], False, None, None),
-            ([4, 13], True, 0.5, 1.5),
-            ([4, 13], False, 0.5, 1.5),
+            ([3, 3], True, None),
+            ([3, 3], False, None),
+            (None, True, [0.5, 1.5]),
+            (None, False, [0.5, 1.5]),
         ]
     )
-    def test_bilinear2d_dynamic_shape(
-        self, output_size, align_corners, scales_h, scales_w
-    ):
+    def test_bilinear2d_dynamic_shape(self, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_bilinear2d.default(
-                    x, output_size, align_corners, scales_h, scales_w
+                return torch.ops.aten.upsample_bilinear2d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         input_specs = [
@@ -212,19 +212,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([7, 7, 7], [3, 3, 3], True, None, None, None),
-            ([7, 7, 7], [3, 3, 3], False, None, None, None),
-            ([7, 7, 7], [3, 7, 10], True, 0.5, 1.0, 1.5),
-            ([7, 7, 7], [3, 7, 10], False, 0.5, 1.0, 1.5),
+            ([7, 7, 7], [3, 3, 3], True, None),
+            ([7, 7, 7], [3, 3, 3], False, None),
+            ([7, 7, 7], None, True, [0.5, 1.0, 1.5]),
+            ([7, 7, 7], None, False, [0.5, 1.0, 1.5]),
         ]
     )
-    def test_trilinear3d(
-        self, input_size, output_size, align_corners, scales_d, scales_h, scales_w
-    ):
+    def test_trilinear3d(self, input_size, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_trilinear3d.default(
-                    x, output_size, align_corners, scales_d, scales_h, scales_w
+                return torch.ops.aten.upsample_trilinear3d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -232,19 +230,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([3, 3, 3], True, None, None, None),
-            ([3, 3, 3], False, None, None, None),
-            ([4, 9, 13], True, 0.5, 1.0, 1.5),
-            ([4, 9, 13], False, 0.5, 1.0, 1.5),
+            ([3, 3, 3], True, None),
+            ([3, 3, 3], False, None),
+            (None, True, [0.5, 1.0, 1.5]),
+            (None, False, [0.5, 1.0, 1.5]),
         ]
     )
-    def test_trilinear3d_dynamic_shape(
-        self, output_size, align_corners, scales_d, scales_h, scales_w
-    ):
+    def test_trilinear3d_dynamic_shape(self, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_trilinear3d.default(
-                    x, output_size, align_corners, scales_d, scales_h, scales_w
+                return torch.ops.aten.upsample_trilinear3d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         input_specs = [
@@ -259,19 +255,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([7, 7], [3, 3], True, None, None),
-            ([7, 7], [3, 3], False, None, None),
-            ([7, 7], [3, 10], True, 0.5, 1.5),
-            ([7, 7], [3, 10], False, 0.5, 1.5),
+            ([7, 7], [3, 3], True, None),
+            ([7, 7], [3, 3], False, None),
+            ([7, 7], None, True, [0.5, 1.5]),
+            ([7, 7], None, False, [0.5, 1.5]),
         ]
     )
-    def test_bicubic2d(
-        self, input_size, output_size, align_corners, scales_h, scales_w
-    ):
+    def test_bicubic2d(self, input_size, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_bicubic2d.default(
-                    x, output_size, align_corners, scales_h, scales_w
+                return torch.ops.aten.upsample_bicubic2d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         inputs = [torch.randn([1, 1] + input_size)]
@@ -279,19 +273,17 @@ def forward(self, x):
 
     @parameterized.expand(
         [
-            ([3, 3], True, None, None),
-            ([3, 3], False, None, None),
-            ([4, 13], True, 0.5, 1.5),
-            ([4, 13], False, 0.5, 1.5),
+            ([3, 3], True, None),
+            ([3, 3], False, None),
+            (None, True, [0.5, 1.5]),
+            (None, False, [0.5, 1.5]),
         ]
     )
-    def test_bicubic2d_dynamic_shape(
-        self, output_size, align_corners, scales_h, scales_w
-    ):
+    def test_bicubic2d_dynamic_shape(self, output_size, align_corners, scale_factors):
         class TestModule(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.aten.upsample_bicubic2d.default(
-                    x, output_size, align_corners, scales_h, scales_w
+                return torch.ops.aten.upsample_bicubic2d.vec(
+                    x, output_size, align_corners, scale_factors
                 )
 
         input_specs = [