Solved by ptrblck in post #2 
              
              
                I assume the second K is smaller than the first one. 
This code should work: 
x = torch.randint(0, 10, (10, 5))
k = 3
kvals, kidx = x.topk(k=k, dim=1)
x.zero_()
x[torch.arange(x.size(0))[:, None], kidx] = 1