Solved by Eta_C in post #2 
              
              
                Hope could help you. 
import torch 

n = 5 
top_k = torch.randint(n, (n,))
print(top_k)

# Step 1, generate a variable number of elements in each row
top_k = torch.nn.functional.one_hot(top_k, num_classes=n)
top_k = 1 - torch.cumsum(top_k, dim=1)  # 
print(top_k)

# Step 2, shuffle each row 
indices…