def predict_multiclass(model, x, *extra_xs):
    """
    Predict a class (int in [0...(n classes - 1)]) for each input (first dim of `x`) using the classifier `model`.

    Args:
        model: multiclass classifier module that outputs "n classes" logits for every input sequence.
        x: input tensor.
        extra_xs: additional inputs of `model`.
    Returns
        Integer tensor of shape `(x.size(0),)` that contains the predicted classes.
    """
    preds = model(x, *extra_xs).cpu()
    pred_classes = preds.argmax(dim=-1)
    return pred_classes