CudnnRNN is not differentiable twice about pytorch-rl HOT 4 CLOSED

def forward(self, inputs):
        x = self.hidden_activation(self.input_layer(inputs))

        # Hidden Layers
        for hidden_layer in self.hidden_layers:
            x = self.hidden_activation(hidden_layer(x))
            
        # GRUCell
        outputs = []
        for seq in range(x.size(1)):
                self.hidden = self.gru(x[:, seq], self.hidden)
                outputs.append(self.hidden)
        x = torch.stack(outputs, 1)
        
        # Output Layer
        action_mean = self.output_layer(x)
        action_log_std = self.a_logstd.expand_as(action_mean)
        action_std = torch.exp(action_log_std)

        return action_mean, action_log_std, action_std

def get_log_prob(self, x, actions):
        self.hidden = self.init_hidden(x.size(0))
        action_mean, action_log_std, action_std = self.forward(x)
        return normal_log_density(actions, action_mean, action_log_std, action_std, is_recurrent=True)

def get_fim(self, x):
        self.hidden = self.init_hidden(x.size(0))
        mean, _, _ = self.forward(x)
        cov_inv = self.a_logstd.data.exp().pow(-2).squeeze(0).repeat(x.size(0))
        param_count = 0
        std_index = 0
        id = 0
        for name, param in self.named_parameters():
            if name == "a_logstd":
                std_id = id
                std_index = param_count
            param_count += param.data.view(-1).shape[0]
            id += 1
        return cov_inv, mean, {'std_id': std_id, 'std_index': std_index}