@albanD  I found another similar way to do it;
Maybe this will help you understand my goal!
    for i in range(batch_size):
        lab    = torch.arange(elementPerClass*label[i], elementPerClass*(label[i]+1), 1, dtype=torch.long)
        if lab.nelement()==0:
            print("Label tensor -- empty weight.")
        else:
            mask[i,lab] = ones