fix(//core/conversion/conversionctx): Check both tensor and eval maps

before throwing conversion warning Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · May 4, 2020 · 2d65ece · 2d65ece
1 parent 78a1c61
commit 2d65ece
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -118,10 +118,13 @@ std::string ConversionCtx::SerializeEngine() {
 
 bool ConversionCtx::CheckLayerAddition(const torch::jit::Node* n) {
     for (auto out : n->outputs()) {
-        auto iter = this->value_tensor_map.find(out);
-        if (iter == this->value_tensor_map.end()) {
-            LOG_WARNING("Node " << util::node_info(n) << " output: " << out->debugName() << " does not have a coresponding output, may potentially indicate a defective converter");
-            return false;
+        auto iter_t = this->value_tensor_map.find(out);
+        if (iter_t == this->value_tensor_map.end()) {
+            auto iter_iv = this->evaluated_value_map.find(out);
+            if (iter_iv == this->evaluated_value_map.end()) {
+                LOG_WARNING("Node " << util::node_info(n) << " output: " << out->debugName() << " does not have a coresponding value or tensor, may potentially indicate a defective evaluator or converter");
+                return false;
+            }
         }
     }
     return true;