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..14e324bbb2 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- pytorch/pytorch#134663" +) class TestChunkDynamicConverter(DispatchTestCase): @parameterized.expand( [