From d0cc035c0f30b218a7f13797a2aa0c7f91e2ef1f Mon Sep 17 00:00:00 2001
From: Eric Junyuan Xie <piiswrong@users.noreply.github.com>
Date: Mon, 12 Sep 2016 15:46:42 -0700
Subject: [PATCH] Fix inplace. Only permit shared when shape and dtype is the
 same (#45)

---
 nnvm/src/pass/plan_memory.cc | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/nnvm/src/pass/plan_memory.cc b/nnvm/src/pass/plan_memory.cc
index 0f0b8e79b9c2..f77e70085fbc 100644
--- a/nnvm/src/pass/plan_memory.cc
+++ b/nnvm/src/pass/plan_memory.cc
@@ -173,7 +173,10 @@ Graph PlanMemory(Graph ret) {
       for (auto& kv : inplace_pairs) {
         uint32_t eid_out = idx.entry_id(nid, kv.second);
         uint32_t eid_in = idx.entry_id(inode.inputs[kv.first]);
-        if (ref_count[eid_in] == 1 && storage[eid_in] != GraphAllocator::kBadStorageID) {
+        if (ref_count[eid_in] == 1 &&
+            storage[eid_in] != GraphAllocator::kBadStorageID &&
+            shape_vec[eid_out].Size() == shape_vec[eid_in].Size() &&
+            dtype_vec[eid_out] == dtype_vec[eid_in]) {
           // inplace optimization
           storage[eid_out] = storage[eid_in];
           ref_count[eid_in] = 0;