Merge pull request #1289 from pytorch/fix_misleading_msg

fix: fix misleading skipping partitioning msg
pytorch · Aug 19, 2022 · 2fc413b · 2fc413b
2 parents abf3d47 + 0591a8d
commit 2fc413b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/core/compiler.cpp b/core/compiler.cpp
@@ -426,7 +426,7 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
       auto outputIsCollection = conversion::OutputIsCollection(g->block());
       if (cfg.partition_info.enabled &&
           (cfg.lower_info.forced_fallback_modules.size() == 0 &&
-           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible)) {
+           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) && !outputIsCollection) {
         LOG_INFO("Skipping partitioning since model is fully supported");
       }