# Buggy code
def forward(self, state, x):
    if state is None:
        h = [
            to_device(self, self.zero_state(x.size(0)))
            for n in six.moves.range(self.n_layers)
        ]
        state = {"h": h}
        if self.typ == "lstm":
            c = [
                to_device(self, self.zero_state(x.size(0)))
                for n in six.moves.range(self.n_layers)
            ]
            state = {"c": c, "h": h}

    h = [None] * self.n_layers
    emb = self.embed(x)
    if self.typ == "lstm":
        c = [None] * self.n_layers
        h[0], c[0] = self.rnn[0](self.dropout[0](emb), (state["h"][0], state["c"][0]))
        for n in six.moves.range(1, self.n_layers):
            h[n], c[n] = self.rnn[n](
                self.dropout[n](h[n - 1]), (state["h"][n], state["c"][n])
            )
        state = {"c": c, "h": h}
    else:
        h[0] = self.rnn[0](self.dropout[0](emb), state["h"][0])
        for n in six.moves.range(1, self.n_layers):
            h[n] = self.rnn[n](self.dropout[n](h[n - 1]), state["h"][n])
        state = {"h": h}
    y = self.lo(self.dropout[-1](h[-1]))
    return state, y


# 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 forward(self, x, t):
    """Reduce framewise loss values

    Args:
        x (torch.Tensor): Input ids. (batch, len)
        t (torch.Tensor): Target ids. (batch, len)

    Returns:
        tuple[torch.Tensor, torch.Tensor]: Tuple of
            the reduced loss value along time (scalar)
            and the number of valid loss values (scalar)
    """
    loss = 0
    count = torch.tensor(0).long()
    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)
    return loss, count.to(loss.device)