Fix cached op

Change-Id: If90c6f0997548ffd5daa67cc18bab7405f24213b
apache · Jun 8, 2019 · 23c2c82 · 23c2c82
1 parent e01c386
commit 23c2c82
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
@@ -595,7 +595,10 @@ inline bool CheckAndInferShape(nnvm::Graph* p_g, mxnet::ShapeVector&& shapes,
     *contain_unknown = false;
   }
   nnvm::Graph& g = *p_g;
-  if (g.attrs.count("shape")) {
+  if (use_inputs) {
+    if (g.attrs.count("shape_inputs") && g.GetAttr<mxnet::ShapeVector>("shape_inputs") == shapes)
+      return true;
+  } else if (g.attrs.count("shape")) {
     const auto& prev_shapes = g.GetAttr<mxnet::ShapeVector>("shape");
     if (prev_shapes.size() == shapes.size()) {
       bool match = true;