How to properly setup parameters when using functional_call and nn.LSTM

cijerezg: <p>I have the following model:</p>
<pre><code class="lang-auto">class Encoder(nn.Module):                                                                                                                                 
      def __init__(self, action_dim, z_dim, skill_length):                                                                                                  
          super().__init__()                                                                                                                                
                                                                                                                                             
          self.lstm = nn.LSTM(action_dim, z_dim, skill_length, batch_first=True)                                                                                                             
                                                                                                                                                                                                                                                               
          self.log_std = nn.Parameter(torch.Tensor(z_dim))                                                                                                                                                                                                                
                                                                                                                                                                                                                                                                                                                    
      def forward(self, skill):                                                                                                                             
          mean, _ = self.lstm(skill)                                                                                                                                                                                                                                                    
          mean = mean[:, -1, :]                                                                                                                             
          std = torch.exp(torch.clamp(self.log_std, min=math.log(epsilon)))                                                                                 
          density = Normal(mean, std)                                                                                                                       
          sample = density.rsample()                                                                                                                        
                                                                                                                                                                                                 
          return sample, density  
</code></pre>
<p>Then, I have the following code to extract and initialize parameters:</p>
<pre><code class="lang-auto">def pars(model):
    params = {}
    for name, param in model.named_parameters():
        if 'std' in name:
            init = torch.nn.init.constant_(param, 0)
        else:
            init = torch.nn.init.orthogonal_(param)
        params[name] = nn.Parameter(init)
    return params
</code></pre>
<p>Then, I’d like to do this:</p>
<pre><code class="lang-auto">model = Encoder(6, 2, 10)
x = torch.rand(25, 10, 6)
params = pars(model)

samp, d = functional_call(model, params, x)

grad = autograd.grad(torch.mean(samp), params.values(), retain_graph=True, allow_unused=True)
</code></pre>
<p>There should be derivative depending on the lstm layer as well as the std layer, but in my case I get <code>None</code> for the lstm layer.</p>
<p>When I run</p>
<pre><code class="lang-auto">samp, d = model(x)
grad = autograd.grad(torch.mean(samp), model.parameters(), retain_graph=True)
</code></pre>
<p>then I get the correct gradient, i.e.g, it depends on the lstm layer as well as the std parameters, instead of getting <code>None</code>.</p>
<p>I am not sure what the problem is. I have computed gradients successfully before using <code>functional_call</code>, but it seems introducing <code>rsample</code> is causing some issues. I’d appreciate any help.</p>
<p>Edit: I found the problem is with the LSTM layer. When I use a linear layer, instead of LSTM, then everything works fine.</p>