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