diff --git a/paddle/fluid/framework/details/memory_optimize_helper.cc b/paddle/fluid/framework/details/memory_optimize_helper.cc index 6126c168ccf4b52ae75648d2dfe48db770639214..db4e805bb692ee44ac50337fae54f8dbfe389e6f 100644 --- a/paddle/fluid/framework/details/memory_optimize_helper.cc +++ b/paddle/fluid/framework/details/memory_optimize_helper.cc @@ -129,7 +129,13 @@ size_t NodeSize(const VarDesc& node) { } size_t NodeSize(ir::Node* n) { - auto* desc = FindVarDescInBlock(n); + VarDesc* desc = nullptr; + // some op do not have block pointer + if (n->inputs[0]->Op() != nullptr) { + desc = FindVarDescInBlock(n); + } else { + desc = n->Var(); + } return NodeSize(*desc); } diff --git a/paddle/fluid/framework/details/memory_optimize_pass.cc b/paddle/fluid/framework/details/memory_optimize_pass.cc index 93d08649db2673b31faad39e5b6d2aea390a8e08..d45a43d851dc7dff1421a4f8e3316a8b1169e979 100644 --- a/paddle/fluid/framework/details/memory_optimize_pass.cc +++ b/paddle/fluid/framework/details/memory_optimize_pass.cc @@ -194,7 +194,8 @@ void MemoryOptimizePass::SubGraphOptimize(OpDesc* op_desc) const { // effect. Because it is a single op in graph. No need to // update the ir nodes. sub_op_desc->Rename(var->Name(), cache->Name()); - if (sub_op_desc->Block()->HasVar(var->Name())) { + if (sub_op_desc->Block() != nullptr && + sub_op_desc->Block()->HasVar(var->Name())) { sub_op_desc->Block()->RemoveVar(var->Name()); } }