Solved by KFrank in post #2 
              
              
                Hi Melike! 

Try: 
>>> import torch
>>> torch.__version__
'1.9.0'
>>> _ = torch.manual_seed (2021)
>>> N = 10
>>> t = torch.randint (N, (2, 3, 5))
>>> i = torch.ones ([N] + list (t.shape)).cumsum (0) - 1
>>> masks = t == i
>>> t
tensor([[[4, 9, 3, 2, 5],
         [6, 2, 0, 5, 0],
         [1, 9, 7, …