Merge pull request #992 from pytorch/anuragd/fix_aten_split

(//core): Added a variant for aten::split
pytorch · May 23, 2022 · 042fa2e · 042fa2e
2 parents 2c787ad + 99d1f32
commit 042fa2e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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);