import torch


def extract_logits(out, task_ids, n_classes, device):
    """
    Extract logits corresponding to task_ids from out.

    Args:
        out: Predictions
        task_ids: Task ids
        n_classes: Number of classes per task
    """
    indices = (
        (torch.arange(n_classes * out.size(0)) % n_classes)
        .reshape(out.size(0), n_classes)
        .to(device=device)
    )
    indices = indices + (task_ids * n_classes).unsqueeze(1)
    return out[torch.arange(out.size(0)).unsqueeze(1), indices]