From d92a6f532f8b40fb867797c89a779dad14fb1210 Mon Sep 17 00:00:00 2001
From: apbose <apbose694@gmail.com>
Date: Wed, 4 Sep 2024 16:02:43 -0700
Subject: [PATCH] removing chunk op, changing the skip test message and remove
 the dynamic shape check of input in validator

---
 .../dynamo/conversion/aten_ops_converters.py  | 28 +------------------
 tests/py/dynamo/conversion/test_chunk_aten.py |  6 ++--
 2 files changed, 5 insertions(+), 29 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index a625c05791..1999bffe4a 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -692,9 +692,7 @@ def aten_ops_softmax(
 
 @dynamo_tensorrt_converter(
     torch.ops.aten.split.Tensor,
-    capability_validator=(
-        has_static_shapes_in_args([0]) and has_static_shapes_in_args([1])
-    ),
+    capability_validator=has_static_shapes_in_args([1]),
     supports_dynamic_shapes=True,
 )
 @dynamo_tensorrt_converter(
@@ -905,30 +903,6 @@ def aten_ops_slice(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.chunk.default)
-@enforce_tensor_types(
-    {
-        0: (TRTTensor,),
-    }
-)
-def aten_ops_chunk(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.slice.chunk(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        args[1],
-        args_bounds_check(args, 2, 0),
-    )
-
-
 @dynamo_tensorrt_converter(torch.ops.aten.cumsum.default, supports_dynamic_shapes=True)
 @enforce_tensor_types(
     {
diff --git a/tests/py/dynamo/conversion/test_chunk_aten.py b/tests/py/dynamo/conversion/test_chunk_aten.py
index eb74854009..eb06c04201 100644
--- a/tests/py/dynamo/conversion/test_chunk_aten.py
+++ b/tests/py/dynamo/conversion/test_chunk_aten.py
@@ -85,8 +85,10 @@ def forward(self, input):
 
 
 #######################Dynamic cases#######################
-#The tests are skipped for now. Will be addressed once https://github.com/pytorch/pytorch/issues/134663 is addressed
-@unittest.skip("Pending aten.split converter. Currently tested by E2E")
+# The tests are skipped for now. Will be addressed once https://github.com/pytorch/pytorch/issues/134663 is addressed
+@unittest.skip(
+    "Pending aten.split dynamic input torch.export guard bug. Issue- https://github.com/pytorch/pytorch/issues/134663"
+)
 class TestChunkDynamicConverter(DispatchTestCase):
     @parameterized.expand(
         [