@@ -320,7 +320,8 @@ std::vector<std::shared_ptr<ExecutorPrepareContext>> Executor::Prepare(
320
320
}
321
321
322
322
void Executor::RunPreparedContext (ExecutorPrepareContext* ctx, Scope* scope,
323
- bool create_local_scope, bool create_vars) {
323
+ bool create_local_scope, bool create_vars,
324
+ bool keep_kids) {
324
325
Scope* local_scope = scope;
325
326
if (create_vars) {
326
327
if (create_local_scope) {
@@ -343,12 +344,20 @@ void Executor::RunPreparedContext(ExecutorPrepareContext* ctx, Scope* scope,
343
344
}
344
345
}
345
346
platform::DeviceContextPool::Instance ().Get (place_)->Wait ();
346
- if (create_vars && create_local_scope ) {
347
+ if (local_scope != scope ) {
347
348
scope->DeleteScope (local_scope);
348
349
} else {
349
- // Delete the local scopes created in operators.
350
- scope->DropKids ();
350
+ if (!keep_kids) {
351
+ // By default, we should delete all kid scopes after run executor because
352
+ // some operators may create local scope when running, such as while_op.
353
+ // But when while_op also create a local executor to run it's sub block,
354
+ // the sub scopes it created should not be dropped immediately, because
355
+ // while_grad_op will use some variables during while_op run, so we need
356
+ // to keep the kids and wait for the outer executor to drop them.
357
+ scope->DropKids ();
358
+ }
351
359
}
360
+
352
361
if (FLAGS_benchmark) {
353
362
VLOG (2 ) << " -------------------------------------------------------" ;
354
363
VLOG (2 ) << " Memory used after deleting local scope: "
0 commit comments