Hi Melike!



 melike:

suggest equivalent Pytorch tensor operation(s) instead.


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, 7, 4]],

        [[8, 0, 0, 2, 4],
         [1, 7, 5, 7, 0],
         [1, 6, 9, 8, 6]]])
>>> masks[0]
tensor([[[False, False, False, False, False],
         [False, False,  True, False,  True],
         [False, False, False, False, False]],

        [[False,  True,  True, False, False],
         [False, False, False, False,  True],
         [False, False, False, False, False]]])
>>> masks[4]
tensor([[[ True, False, False, False, False],
         [False, False, False, False, False],
         [False, False, False, False,  True]],

        [[False, False, False, False,  True],
         [False, False, False, False, False],
         [False, False, False, False, False]]])

Best.
K. Frank