How am I supposed to cache Tensors that require grad but arent learnable Module params that are replaced by a different Tensor several times each forward pass?

username: marvosyntactical
cooked: <p>I’m not 100% sure if the question actually addresses my problem so let me restate it:<br>
I am trying to get the Jacobian of an input batch to an RNN(nn.Module), but after calling output.grad() the input.grad is still None. I closed in on what I think the problem is with these strategies:</p>
<ol>
<li>I double checked the input and all of the model parameters require grad</li>
<li>I went from the outputs (which have a grad_fn) back through the model to find the last param whose param.grad is None.</li>
</ol>
<p>This last gradless param is the net.activs parameter which gets replaced on every RNN step with a new nn.Parameter instance containing the activation values:</p>
<pre><code class="lang-auto">class RNN(nn.Module):
def __init__(self,many,args):
    # ...
    self.requires_grad = None # set this appropriately before forward pass to know how to 
def reset(self, batch_size=None):
        """reset internal batch activation tensors which arent learnable model params"""
        if batch_size is None:
            batch_size = self.batch_size
        else:
            self.batch_size = batch_size
        if self.n_hidden &gt; 0:
            self.activs = nn.Parameter(torch.zeros(batch_size, self.n_hidden, dtype=self.dtype))
        else:
            self.activs = None
        self.outputs = nn.Parameter(torch.zeros(batch_size, self.n_outputs, dtype=self.dtype))
      
def forward(self, inputs):
        inputs = torch.tensor(inputs, dtype=self.dtype)
        activs_for_output = self.activs
        if self.n_hidden &gt; 0:
            # recurrent loopdeloop
            for _ in range(self.n_internal_steps):
                self.activs = nn.Parameter(self.activation(
                    self.hidden_responses * (
                        self.input_to_hidden(inputs) +
                        self.hidden_to_hidden(self.activs) +
                        self.output_to_hidden(self.outputs)
                        ) + self.hidden_biases)).requires_grad_(self.requires_grad)
            if self.use_current_activs:
                activs_for_output = self.activs 
        output_inputs = (self.input_to_output(inputs) +
                            self.output_to_output(self.outputs))
        if self.n_hidden &gt; 0:
            output_inputs += self.hidden_to_output(activs_for_output)
        self.outputs = nn.Parameter(self.activation(
            self.output_responses * output_inputs + self.output_biases)).requires_grad_(self.requires_grad)
        return self.outputs

def main():
    net = RNN()
    net.requires_grad_(True)
    net.requires_grad = True

    # try to get jacobian of probe batch
    inputs = torch.rand(batch_size, features).requires_grad_(True)
    outputs = net(inputs)
    outputs.backward(torch.ones_like(outputs))
    
    # forward
    assert states.requires_grad == True
    assert net.activs.requires_grad == True
    assert net.outputs.requires_grad == True
    assert actions.requires_grad == True

    # backward
    assert net.outputs.grad is not None # net.outputs == outputs
    assert net.activs.grad is not None # &lt;---- the error if I remove all nn.Parameter() wraps above
    assert states.grad is not None # &lt;---- my error
    
    jacobian = states.grad.clone().detach().view(batch_size, -1).cpu().numpy()
</code></pre>
<p>I think the problem is that I replace the net.activs parameter with a new instance several times during net.forward, and refer to it on successive forward passes. Should I not set this parameter as an attribute of the model? If yes, how else should I cache it?</p>
<p>What I’ve tried:</p>
<ol>
<li>See the alternative error in above code: try not to use nn.Parameters, leading to the output having no grad</li>
<li>replace self.activs(/outputs) assignments by self.activs[:], yielding: RuntimeError: leaf variable has been moved into the graph interior</li>
</ol>
<p>I will try caching these Tensors somewhere else now.</p>

username: marvosyntactical
cooked: <p>Stopped suppressing warnings and got <em>UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won’t be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor<br>
instead. See <a href="http://github.com/pytorch/pytorch/pull/30531" rel="noopener nofollow ugc">github.com/pytorch/pytorch/pull/30531</a> for more informations.</em><br>
which Im now reading</p>

username: marvosyntactical
cooked: <p>I solved my problem by:</p>
<ol>
<li>Not making activs and outputs nn.Parameters</li>
<li>Not assigning them as model attributes. Instead I added them as optional key word arguments in the forward method and returned the activations as well.</li>
</ol>