Register_full_backward_hook alters behaviour of forward hook

schaefertim: <p>I’d like to update <a href="https://github.com/f-dangel/backpack/" rel="noopener nofollow ugc">BackPACK</a> from <code>register_backward_hook</code> to <code>register_full_backward_hook</code>. However, this changes the forward hook registered with <code>register_forward_hook</code>.<br>
<strong>Expected/old behaviour:</strong> In the forward hook, the argument <code>input[0]</code> of the succeeding module is identical with the argument <code>output</code> of the preceding module.<br>
<strong>New behaviour:</strong> These objects are different now. This can be seen from the ids in the minimum example.<br>
<strong>What I am looking for:</strong> Is this the intended behaviour?<br>
If no, will it be changed to the old behaviour?<br>
If yes, how can I achieve the following:<br>
Based on the minimum example, I calculate a quantity <code>a_loss</code> in the backward hook on <code>_loss</code>. How can I access the quantity <code>a_loss</code> in the backward hook on <code>_linear2</code> to calculate <code>a_linear2</code>, and so on?<br>
If it is not clear what I mean, I can follow up with a separate example.<br>
<strong>Minimum example</strong></p>
<pre><code class="lang-auto">from typing import List, Tuple

import torch
from torch import Tensor
from torch.nn import CrossEntropyLoss, Linear, Module, Sequential


def _forward_hook(module: Module, input: Tuple[Tensor], output: Tensor) -&gt; None:
    print(
        f"{type(module)}: id of input, input[0], and output: "
        f"{id(input)}, {id(input[0])}, {id(output)}"
    )


def _backward_hook(module: Module, g_inp: Tuple[Tensor], g_out: Tuple[Tensor]) -&gt; None:
    pass


def _conduct_experiment(installed_hook: str) -&gt; None:
    batch_size: int = 2
    num_class: int = 3
    _module: Sequential = Sequential(
        Linear(in_features=num_class, out_features=num_class),
        Linear(in_features=num_class, out_features=num_class),
    )
    _input: Tensor = torch.rand(batch_size, num_class, requires_grad=True)
    _target: Tensor = torch.empty(batch_size, dtype=torch.long).random_(num_class)
    _loss_module: CrossEntropyLoss = CrossEntropyLoss()
    list_modules: List[Module] = list(_module.modules()) + [_loss_module]
    print(f"Register hooks with {installed_hook}.")
    for module_child in list_modules:
        module_child.register_forward_hook(_forward_hook)
        getattr(module_child, installed_hook)(_backward_hook)

    print("Forward pass.")
    _loss_module(_module(_input), _target)


_conduct_experiment("register_backward_hook")
print("\n\n")
_conduct_experiment("register_full_backward_hook")
</code></pre>
<p><strong>Output</strong></p>
<pre><code class="lang-auto">Register hooks with register_backward_hook.
Forward pass.
&lt;class 'torch.nn.modules.linear.Linear'&gt;: id of input, input[0], and output: 140708990519632, 140707567033312, 140708989826944
&lt;class 'torch.nn.modules.linear.Linear'&gt;: id of input, input[0], and output: 140708990519632, 140708989826944, 140708989843568
&lt;class 'torch.nn.modules.container.Sequential'&gt;: id of input, input[0], and output: 140708989964432, 140707567033312, 140708989843568
&lt;class 'torch.nn.modules.loss.CrossEntropyLoss'&gt;: id of input, input[0], and output: 140708990430336, 140708989843568, 140707010881248



Register hooks with register_full_backward_hook.
Forward pass.
&lt;class 'torch.nn.modules.linear.Linear'&gt;: id of input, input[0], and output: 140708990389904, 140707010514688, 140707010514848
&lt;class 'torch.nn.modules.linear.Linear'&gt;: id of input, input[0], and output: 140708990443088, 140707010514848, 140707010531728
&lt;class 'torch.nn.modules.container.Sequential'&gt;: id of input, input[0], and output: 140708989964432, 140707010514448, 140707010531808
&lt;class 'torch.nn.modules.loss.CrossEntropyLoss'&gt;: id of input, input[0], and output: 140708989841728, 140707010531808, 140707010532128
</code></pre>

schaefertim: <p>ping <a class="mention" href="/u/alband">@albanD</a> : Maybe you have an idea whether that is the intended behaviour?<br>
I’m grateful for every bit of information.</p>

albanD: <p>Ho sorry, I missed this.</p>
<blockquote>
<p>Is this the intended behaviour?</p>
</blockquote>
<p>It is expected yes that they are different python objects. That being said, they are exactly the same Tensor (same content).</p>
<p>So they do have different “id” but you can use one or the other as they contain the exact same thing <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>