It seems logical indexing does not support broadcasting, is this correct?
I’ve tried following: Let’s say we have a batch of masks and a batch of images, and we want to extract the masked pixels:
mask = torch.randn((10, 1, 32, 32)) > 0
img = torch.rand((10, 3, 32, 32))
img[mask] # error
(We can of course use workarounds using a product of mask and image where broadcasting seems to work.)