def compute_loss(logits, target, criterion):
    """
    logits shape: [seq-len, batch-size, output_dim]
    target shape: [batch-size, seq-len]

    """

    loss = criterion(logits, target)

    return loss