Hi,i implemented my own custom LSTMCell based on [https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py 9],
but during back-propagation i get nan values (after two or three iterations).To be more specific my net is consisted of CNN (Alexnet) + CustomRNN + Log_Softmax and is trained with CTC loss.As far as my custom LSTM is concerned, it is an implementation of Differential RNN https://arxiv.org/abs/1504.06678 1. Below are some snippets of my code.
LSTMCell:
Screenshot from 2019-08-20 16:33:54.png1059×819 112 KB

Model forward:
def forward(self, x):
        LSTMState = namedtuple('LSTMState', ['hx', 'cx', 'dc'])

        batch_size, timesteps, C, H, W = x.size()
        c_in = x.view(batch_size * timesteps, C, H, W)
        c_out = self.cnn(c_in)
        c_out = c_out.view(-1, batch_size, 4096)

        h1 = torch.zeros(batch_size, self.hidden_size).cuda(0)
        h2 = torch.zeros(batch_size, self.hidden_size).cuda(0)
        h3 = torch.zeros(batch_size, self.hidden_size).cuda(0)

        states = [[LSTMState(h1, h2, h3)
                   for _ in range(2)]
                  for _ in range(self.num_layers)]
        r_out, out_state = self.rnn(c_out, states)
        custom_state = double_flatten_states(out_state)
        r_out2 = self.last_linear(r_out)
        return (r_out2)

Thanks in advance