Solved by InnovArul in post #2 
              
              
                In the __init__() function, try making the hidden and cell state as nn.Parameters. 
self.hidden = nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
self.cell_state = nn.Parameter(torch.zeros(1, 1, self.hidden_dim))


And no need to init_hidden() for every batch, as they are learned.