Hi all, I have a 5-dimensional tensor of masks and I want to create a 4-D tensor by masking a 4-D input with each of these masks. So far I could come up with a for loop on unbinded masks, but I’d appreciate if someone could suggest a way to eliminate the for loop since the number of masks is high (hundreds). Here is an example,
>>> import torch

>>> _ = torch.manual_seed(16)
>>> num_masks = 2
>>> inp = torch.randint(1, 5, (3, 4), dtype=torch.float)
>>> outp = torch.zeros(inp.shape)
>>> masks = torch.randn((num_masks,) + inp.shape) > 0
>>> for m in torch.unbind(masks, dim=0):
>>> 	outp[m] = inp[m]


>>> inp
tensor([[2., 3., 2., 2.],
        [2., 4., 1., 2.],
        [4., 1., 2., 1.]])

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

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

>>> outp
tensor([[0., 3., 2., 2.],
        [2., 4., 0., 2.],
        [4., 1., 2., 1.]])