Skip to content

Commit

Permalink
Fix inplace. Only permit shared when shape and dtype is the same (apa…
Browse files Browse the repository at this point in the history
  • Loading branch information
piiswrong authored and tqchen committed May 29, 2018
1 parent 2c0b79a commit d0cc035
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion nnvm/src/pass/plan_memory.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit d0cc035

Please sign in to comment.