From f36a8c6ce98e4903ed3ac37e053e4fc0fdc4c38a Mon Sep 17 00:00:00 2001 From: Anurag Dixit Date: Wed, 20 Apr 2022 19:08:54 -0700 Subject: [PATCH 1/2] (//core): Added a variant for aten::split Signed-off-by: Anurag Dixit --- core/conversion/converters/impl/select.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp index 9c615b63d0..a151cc402d 100644 --- a/core/conversion/converters/impl/select.cpp +++ b/core/conversion/converters/impl/select.cpp @@ -350,6 +350,12 @@ auto select_registrations TORCHTRT_UNUSED = LOG_DEBUG("Converted split op into a list of IValues"); return true; }}) + .pattern({"aten::split.sizes(Tensor(a -> *) self, int[] split_size, int dim=0) -> (Tensor[])", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + add_split(ctx, n, args, true); + LOG_DEBUG("Converted split op into a list of IValues"); + return true; + }}) .pattern({"aten::split.Tensor(Tensor(a) self, int split_size, int dim=0) -> (Tensor[])", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { add_split(ctx, n, args, false, false); From 99d1f328d6371e96b53e7c67ad57380364493eff Mon Sep 17 00:00:00 2001 From: Anurag Dixit Date: Mon, 9 May 2022 18:28:48 -0700 Subject: [PATCH 2/2] fix: Update the call Signed-off-by: Anurag Dixit --- core/conversion/converters/impl/select.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp index a151cc402d..27b29b4195 100644 --- a/core/conversion/converters/impl/select.cpp +++ b/core/conversion/converters/impl/select.cpp @@ -352,7 +352,7 @@ auto select_registrations TORCHTRT_UNUSED = }}) .pattern({"aten::split.sizes(Tensor(a -> *) self, int[] split_size, int dim=0) -> (Tensor[])", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - add_split(ctx, n, args, true); + add_split(ctx, n, args, true, false); LOG_DEBUG("Converted split op into a list of IValues"); return true; }})