A training loop that looked like this would be nicer.
for batch, i in enumerate(range(0, train_data.size(0) - 1, args.bptt)):
    data, targets = get_batch(train_data, i)
    model.zero_grad()
    output = model(data)
    loss = criterion(output.view(-1, ntokens), targets)
    loss.backward()

But for that model.forward would have to look something like this
 def forward(self, input):
     ...
     if self.hidden is None or batch_size_doesnt_match_hidden_size:
         self.hidden = self.init_hidden()
     output, self.hidden = self.rnn(emb, self.hidden.detach())
     ...
     return just_the_output

How about adding a keep_hidden=True parameter to the LSTM class which would ask it to manage its hidden state internally?