fix: [collection] skip ListConstruct in condition

Signed-off-by: inocsin <[email protected]>
pytorch · Apr 8, 2022 · 6391ae3 · 6391ae3
1 parent 47455d1
commit 6391ae3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/core/conversion/conversion.cpp b/core/conversion/conversion.cpp
@@ -498,7 +498,9 @@ std::string ConvertBlockToEngine(
 std::unordered_map<c10::OperatorName, std::string> GetUnsupportedOpsInBlock(const torch::jit::Block* b) {
   std::unordered_map<c10::OperatorName, std::string> unsupported_ops;
   for (const auto n : b->nodes()) {
-    if (n->kind() != torch::jit::prim::Loop && n->kind() != torch::jit::prim::If && !OpSupported(n) && n->kind() != torch::jit::prim::TupleConstruct) {
+    if (n->kind() != torch::jit::prim::Loop && n->kind() != torch::jit::prim::If &&
+       !OpSupported(n) && n->kind() != torch::jit::prim::TupleConstruct &&
+       n->kind() != torch::jit::prim::ListConstruct) {
       auto schema = n->maybeSchema();
       TORCHTRT_CHECK(
           schema,