import torch


def get_topk_class(probability, top_k):
    """
        select top k class according probability for multi-label classification
        :param probability: Tensor of (batch_size, class_size), probability
        :param topk: int, number of selected classes
        :returns: Tensor of long (batch_size, class_size), whether class is selected, 1 means class selected, 0 means class unselected
    """
    # batch_size * class_size
    device = probability.device
    batch_size, class_size = probability.size()
    # batch_size * k
    label_index = probability.topk(top_k, dim=-1, sorted=False)[1]
    batch_index = torch.arange(batch_size).unsqueeze(-1).expand(-1, top_k)
    # batch_size * class_size
    label = torch.zeros(batch_size, class_size).long()
    batch_index = batch_index.to(device)
    label = label.to(device)
    label[batch_index, label_index] = 1
    return label