Merge pull request tqchen#9 from ZihengJiang/bugfix

fix shape infer bug
dlsys-course · Oct 28, 2016 · 05ba599 · 05ba599
2 parents 78dedfd + c4fed46
commit 05ba599
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/session.cc b/src/session.cc
@@ -294,10 +294,10 @@ void TorchExecutor::SetupShapeDType(
       CHECK(state != nullptr);
       CHECK(state->initialized())
           << "Attempt to execute a graph un-initialized Variable";
-      if (node_shape_->at(nid) != state->blob.shape) {
+      if (node_shape_->at(idx.entry_id(nid, 0)) != state->blob.shape) {
         need_redo_infer = true; break;
       }
-      if (node_dtype_->at(nid) != state->blob.dtype) {
+      if (node_dtype_->at(idx.entry_id(nid, 0)) != state->blob.dtype) {
         need_redo_infer = true; break;
       }
     }