Thanks for identifying the issue. It has been solved by passing the batch size to the initHidden function
    def initHidden(self, batch_size):
        return torch.zeros(batch_size, self.hidden_size).to(device)