From: Yinghai Lu <yinghai@fb.com>
Date: Wed, 5 Dec 2018 05:50:41 +0000 (-0800)
Subject: Add output info when doing onnxGetBackendCompatibility (#14784)
X-Git-Tag: accepted/tizen/6.5/unified/20211028.231830~2463
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=ecc17fe3dd03962f36a989659336e42de86a38ca;p=platform%2Fupstream%2Fpytorch.git

Add output info when doing onnxGetBackendCompatibility (#14784)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14784

TSIA. To give more complete info to `onnxGetBackendCompatibility`.

Reviewed By: bertmaher, rdzhabarov

Differential Revision: D13331989

fbshipit-source-id: 1064b93f7f474788f736e6f0c893dae915c6fb99
---

diff --git a/caffe2/opt/onnxifi_transformer.cc b/caffe2/opt/onnxifi_transformer.cc
index 6338b68..b52833f 100644
--- a/caffe2/opt/onnxifi_transformer.cc
+++ b/caffe2/opt/onnxifi_transformer.cc
@@ -419,15 +419,23 @@ void OnnxifiTransformer::Transform(
         onnx_model.mutable_graph()->add_node()->CopyFrom(n);
       }
 
-      // Add input shape info
-      std::vector<std::string> input_tmp;
+      // Add input/output shape info
+      std::vector<std::string> io_tmp;
       for (const auto& op_input : op.input()) {
-        input_tmp.emplace_back(op_input);
+        io_tmp.emplace_back(op_input);
       }
-      auto io_vec = ConvertToValueInfo(input_tmp, shape_hints);
+      auto io_vec = ConvertToValueInfo(io_tmp, shape_hints);
       for (const auto& i : io_vec) {
         onnx_model.mutable_graph()->add_input()->CopyFrom(i);
       }
+      io_tmp.clear();
+      for (const auto& op_output : op.output()) {
+        io_tmp.emplace_back(op_output);
+      }
+      io_vec = ConvertToValueInfo(io_tmp, shape_hints);
+      for (const auto& i : io_vec) {
+        onnx_model.mutable_graph()->add_output()->CopyFrom(i);
+      }
 
       std::string onnx_model_str;
       onnx_model.SerializeToString(&onnx_model_str);