I couldn’t spot anything abnormal… Below is my model
  class PredRNN(nn.Module):
  def __init__(self, rnn_type, input_size, hidden_size, output_size, n_layers=1, dropout=0.0,
             bidirectional=False):
    super(PredRNN, self).__init__()
    self.input_size = input_size
    self.hidden_size = hidden_size
    self.output_size = output_size
    self.n_layers = n_layers
    self.rnn_type = rnn_type
    self.dropout = dropout
    if bidirectional:
        self.multiply = 2
    else:
        self.multiply = 1
    if self.rnn_type == 'lstm':
        self.rnn = nn.LSTM(self.input_size, self.hidden_size, self.n_layers, dropout=self.dropout)
    elif self.rnn_type == 'gru':
        self.rnn = nn.GRU(self.input_size, self.hidden_size, self.n_layers, dropout=self.dropout)
    else:
        self.rnn = nn.LSTM(self.input_size, self.hidden_size, self.n_layers, dropout=self.dropout)

    self.rnn2out = nn.Linear(self.hidden_size * self.multiply, self.output_size)
    self.hidden = self.init_hidden()

def forward(self, inputs):
    output, self.hidden = self.rnn(inputs, self.hidden)
    output = self.rnn2out(output.view(len(inputs), -1))
    output = F.log_softmax(output)
    return output

def init_hidden(self):
    if self.rnn_type == 'lstm':
        h0 = Variable(torch.zeros(self.n_layers * self.multiply, 1, self.hidden_size))
        c0 = Variable(torch.zeros(self.n_layers * self.multiply, 1, self.hidden_size))
        if USE_CUDA:
            hidden = (h0.cuda(), c0.cuda())
        else:
            hidden = (h0, c0)
    elif self.rnn_type == 'gru':
        hidden = Variable(torch.zeros(self.n_layers * self.multiply, 1, self.hidden_size))
        if USE_CUDA:
            hidden = hidden.cuda()
    else:
        h0 = Variable(torch.zeros(self.n_layers * self.multiply, 1, self.hidden_size))
        c0 = Variable(torch.zeros(self.n_layers * self.multiply, 1, self.hidden_size))
        if USE_CUDA:
            hidden = (h0.cuda(), c0.cuda())
        else:
            hidden = (h0, c0)

    return hidden