aten::cat rebase changes

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -53,19 +53,19 @@ def aten_ops_batch_norm(
 
 @dynamo_tensorrt_converter(torch.ops.aten.cat.default)
 def aten_ops_cat(
-    network: TRTNetwork,
+    ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return impl.cat.cat(
-        network,
+        ctx,
         target,
         SourceIR.ATEN,
         name,
-        tensors=args[0],
-        dim=args_bounds_check(args, 2, 1),
+        input=args[0],
+        dim=args_bounds_check(args, 1, 0),
     )
 
 

py/torch_tensorrt/dynamo/conversion/impl/cat.py

@@ -4,22 +4,23 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
 
 def cat(
-    network: TRTNetwork,
+    ctx: ConversionContext,
     target: Target,
     source_ir: Optional[SourceIR],
     name: str,
-    input: TRTNetwork,
+    input: Union[TRTTensor, Sequence[TRTTensor]],
     dim: int,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    if any(not isinstance(t, TRTTensor) for t in input):  # type: ignore[union-attr]
-        raise RuntimeError(
-            f"cat received inputs {input} that is not part " "of the TensorRT region!"
-        )
-    concat_layer = network.add_concatenation(input)
+    for each_input in input:
+        if(not isinstance(each_input, TRTTensor)):
+            each_input = get_trt_tensor(each_input)
+    concat_layer = ctx.net.add_concatenation(input)
     if dim < 0:
         dim = len(input[0].shape) + dim