Add split_with_sizes converter (pytorch#71953)

Summary:
Pull Request resolved: pytorch#71953

Add converter for split_with_sizes

Reviewed By: yinghai

Differential Revision: D33829024

fbshipit-source-id: 50de383797a347ef7afecfbda80b2c84e244e404

Author: Shirong Wu

test/fx2trt/converters/acc_op/test_split.py

@@ -32,5 +32,23 @@ def forward(self, x):
             test_explicit_batch_dim=False,
         )
 
+    @parameterized.expand(
+        [
+            ("split_with_size", [2, 3, 5], 1),
+        ]
+    )
+    def test_split_with_size(self, _, split_size, dim):
+        class Split(nn.Module):
+            def forward(self, x):
+                return x.split_with_sizes(split_size, dim)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(
+            Split(),
+            inputs,
+            expected_ops={acc_ops.slice_tensor},
+            test_explicit_batch_dim=False,
+        )
+
 if __name__ == '__main__':
     run_tests()

torch/fx/experimental/fx2trt/converters/acc_ops_converters.py

@@ -1661,7 +1661,7 @@ def slice_to_trt_params(py_slice, dim_size):
         size = math.ceil((stop - start) * 1.0 / stride)
         return start, size, stride
 
-    if not isinstance(slices, tuple):
+    if not isinstance(slices, tuple) and not isinstance(slices, list):
         slices = (slices,)
 
     if network.has_implicit_batch_dimension:

torch/fx/experimental/fx_acc/acc_ops.py

@@ -1277,6 +1277,14 @@ def linalg_norm(*, input, ord, dim, keepdim):
         ("dim", "dim"),
     ],
 )
+@register_custom_acc_mapper_fn(
+    op_and_target=("call_method", "split_with_sizes"),
+    arg_replacement_tuples=[
+        ("tensor", "input"),
+        ("split_sizes", "split_size_or_sections"),
+        ("dim", "dim"),
+    ],
+)
 @register_custom_acc_mapper_fn(
     op_and_target=("call_function", torch.split),
     arg_replacement_tuples=[