# Buggy code
def update_core(self):
    # When we pass one iterator and optimizer to StandardUpdater.__init__,
    # they are automatically named 'main'.
    train_iter = self.get_iterator("main")
    optimizer = self.get_optimizer("main")
    # Progress the dataset iterator for sentences at each iteration.
    batch = train_iter.__next__()
    x, t = concat_examples(batch, device=self.device, padding=(0, -100))
    # Concatenate the token IDs to matrices and send them to the device
    # self.converter does this job
    # (it is chainer.dataset.concat_examples by default)
    loss = 0
    count = 0
    state = None
    batch_size, sequence_length = x.shape
    for i in six.moves.range(sequence_length):
        # Compute the loss at this time step and accumulate it
        state, loss_batch = self.model(state, x[:, i], t[:, i])
        non_zeros = torch.sum(x[:, i] != 0, dtype=torch.float)
        loss += loss_batch * non_zeros
        count += int(non_zeros)

    loss = loss.mean()
    reporter.report({"loss": float(loss.detach())}, optimizer.target)
    reporter.report({"count": count}, optimizer.target)
    # update
    self.model.zero_grad()  # Clear the parameter gradients
    loss.backward()  # Backprop
    if self.gradclip is not None:
        nn.utils.clip_grad_norm_(self.model.parameters(), self.gradclip)
    optimizer.step()  # Update the parameters


# Traceback (RuntimeError)
Exception in main training loop: CUDA out of memory. Tried to allocate 19.62 MiB (GPU 2; 22.38 GiB total capacity; 21.01 GiB already allocated; 10.06 MiB free; 634.29 MiB cached)
Traceback (most recent call last):
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run
update()
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/updaters/standard_updater.py", line 175, in update
self.update_core()
File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 277, in update_core
loss.backward()  # Backprop
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 102, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 90, in backward
allow_unreachable=True)  # allow_unreachable flag
Will finalize trainer extensions and updater before reraising the exception.
Traceback (most recent call last):
File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/bin/lm_train.py", line 156, in <module>
main(sys.argv[1:])
File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/bin/lm_train.py", line 150, in main
train(args)
File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 416, in train
trainer.run()
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 349, in run
six.reraise(*exc_info)
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/six.py", line 693, in reraise
raise value
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/trainer.py", line 316, in run
update()
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/chainer/training/updaters/standard_updater.py", line 175, in update
self.update_core()
File "/mnt/cfs1/asr/users/fanlu/espnet/espnet/lm/pytorch_backend/lm.py", line 277, in update_core
loss.backward()  # Backprop
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 102, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/mnt/cfs1/asr/users/fanlu/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 90, in backward
allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 19.62 MiB (GPU 2; 22.38 GiB total capacity; 21.01 GiB already allocated; 10.06 MiB free; 634.29 MiB cached)
# Accounting: time=746 threads=1
# Finished at Sun Aug 11 14:58:25 CST 2019 with status 1

# Fixed code
def update_core(self):
    # When we pass one iterator and optimizer to StandardUpdater.__init__,
    # they are automatically named 'main'.
    train_iter = self.get_iterator("main")
    optimizer = self.get_optimizer("main")
    # Progress the dataset iterator for sentences at each iteration.
    batch = train_iter.__next__()
    # Concatenate the token IDs to matrices and send them to the device
    # self.converter does this job
    # (it is chainer.dataset.concat_examples by default)
    x, t = concat_examples(batch, device=self.device, padding=(0, -100))
    loss, count = self.model(x, t)
    loss = loss.sum()
    count = count.sum()
    reporter.report({"loss": float(loss.detach() / count)}, optimizer.target)
    reporter.report({"count": int(count)}, optimizer.target)
    # update
    self.model.zero_grad()  # Clear the parameter gradients
    loss.backward()  # Backprop
    if self.gradclip is not None:
        nn.utils.clip_grad_norm_(self.model.parameters(), self.gradclip)
    optimizer.step()  # Update the parameters