feat: Enable prim::DictConstruct to fallback without conversion check…

… error Signed-off-by: Dheeraj Peri <[email protected]>
pytorch · Mar 17, 2022 · 01d98c7 · 01d98c7
1 parent 11bcb98
commit 01d98c7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/core/conversion/conversion.cpp b/core/conversion/conversion.cpp
@@ -487,7 +487,8 @@ 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)) {
+    if (n->kind() != torch::jit::prim::Loop && n->kind() != torch::jit::prim::If && !OpSupported(n) &&
+        n->kind() != torch::jit::prim::DictConstruct) {
       auto schema = n->maybeSchema();
       TORCHTRT_CHECK(
           schema,