diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp index 9c615b63d0..27b29b4195 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, false); + 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);