<p dir="auto">In PR <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="331398278" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/8354" data-hovercard-type="pull_request" data-hovercard-url="/pytorch/pytorch/pull/8354/hovercard" href="https://github.com/pytorch/pytorch/pull/8354">#8354</a> we noted that the autograd engine always clears its final callbacks on execute():</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="auto Engine::execute(const edge_list&amp; roots,
                     const variable_list&amp; inputs,
                     bool keep_graph,
                     bool create_graph,
                     const edge_list&amp; outputs) -&gt; variable_list {
  std::call_once(start_threads_flag, &amp;Engine::start_threads, this);

  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
  validate_outputs(roots, const_cast&lt;variable_list&amp;&gt;(inputs), [](const std::string&amp; msg) {
    return msg;
  });

  // Callbacks are only valid for the duration of this run and should always be cleared
  ClearCallbacks _cb_guard(final_callbacks, post_callbacks_lock);"><pre class="notranslate"><code class="notranslate">auto Engine::execute(const edge_list&amp; roots,
                     const variable_list&amp; inputs,
                     bool keep_graph,
                     bool create_graph,
                     const edge_list&amp; outputs) -&gt; variable_list {
  std::call_once(start_threads_flag, &amp;Engine::start_threads, this);

  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
  validate_outputs(roots, const_cast&lt;variable_list&amp;&gt;(inputs), [](const std::string&amp; msg) {
    return msg;
  });

  // Callbacks are only valid for the duration of this run and should always be cleared
  ClearCallbacks _cb_guard(final_callbacks, post_callbacks_lock);
</code></pre></div>
<p dir="auto">Which suggests that if execute() is called twice during a run of backward() then callbacks already enqueued from the first or outermost run will be destroyed. This is clearly unexpected.</p>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/apaszke/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/apaszke">@apaszke</a> and I think we could address this by associating final callbacks with backwards graphs and calling them as each graph finishes. We may also want to throw an error if a final callback is created outside of a backward graph it can be associated with.</p>
<p dir="auto">Edit: code formatting.</p>
<p dir="auto">Currently, <code class="notranslate">final_callbacks_</code> of an autograd engine is a vector shared by all backward passes on this engine.  The autograd engine <a href="https://github.com/pytorch/pytorch/blob/936632b120aafd2d12233c1096d9000fe9554b77/torch/csrc/autograd/engine.cpp#L666">clears</a> final callbacks at the beginning of the backward call. With reentrant backward calls, the final callback enqueued by the original backward call will be cleared by the reentrant ones (learnt from <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pritamdamania87/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pritamdamania87">@pritamdamania87</a>). I was a little surprised on this API behavior as I queued a callback but it never fires. Would it be better if we attach these final callbacks to each graph task? For example, we could use a thread local stack to hold the current graph task, and then push and pop accordingly similar to what we did for <code class="notranslate">current_depth</code>. Thoughts?</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ezyang/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ezyang">@ezyang</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ssnl/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ssnl">@ssnl</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/albanD/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/albanD">@albanD</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zou3519/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zou3519">@zou3519</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gqchen/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gqchen">@gqchen</a></p>