def compute_loss(pred, gold, ntokens, criterion):
    """
    Parameters
    ----------
    pred      : 3d tensor (batch_size, seq_len, n_vocab)
        this is the output of the transformer model
    gold      : 2d tenasor of int (batch_size, seq_len)
        target sequence
    ntokens   : int
        number of valid tokens
    criterion : function(2d tensor, 1d long tensor)
        a loss function. 1st arg is log probability distribution.
        2nd arg is raw gold data. the index of first argument corresponds
        to the value of 2nd arg.
    """
    # flatten pred and gold
    pred = pred.reshape(-1, pred.shape[2])
    gold = gold.reshape(-1)
    loss = criterion(pred, gold) / ntokens
    return loss