The following toy example illustrates the issue: the masked outputs for two different data points in the same batch are identical, despite the data points themselves being distinct.
import torch
x = torch.rand((2,10,3))
masks=torch.tensor([[1,2,4],[1,2,4]])
print(f'x.shape :{x.shape}')
print(f'x: {x}')
print(f'masks.shape :{masks.shape}')
print(f'masks: {masks}')
print(apply_masks(x, masks,concat=True))
x.shape :torch.Size([2, 10, 3])
x: tensor([[[0.2358, 0.6246, 0.6114],
[0.0488, 0.4238, 0.9856],
[0.6111, 0.5021, 0.0529],
[0.2561, 0.5448, 0.7761],
[0.9557, 0.6535, 0.0349],
[0.2390, 0.1336, 0.6073],
[0.6863, 0.7024, 0.9983],
[0.4782, 0.6668, 0.3349],
[0.3228, 0.5506, 0.3206],
[0.7068, 0.4859, 0.3012]],
[[0.7707, 0.7672, 0.0812],
[0.8873, 0.3862, 0.6775],
[0.8626, 0.9021, 0.9553],
[0.5033, 0.5018, 0.9031],
[0.4959, 0.1582, 0.1944],
[0.7786, 0.9554, 0.3488],
[0.5261, 0.1865, 0.0872],
[0.4738, 0.1782, 0.4749],
[0.3570, 0.5354, 0.6668],
[0.7930, 0.2098, 0.4274]]])
masks.shape :torch.Size([2, 3])
masks: tensor([[1, 2, 4],
[1, 2, 4]])
tensor([[[0.0488, 0.4238, 0.9856],
[0.6111, 0.5021, 0.0529],
[0.9557, 0.6535, 0.0349]],
[[0.0488, 0.4238, 0.9856],
[0.6111, 0.5021, 0.0529],
[0.9557, 0.6535, 0.0349]]])
The following toy example illustrates the issue: the masked outputs for two different data points in the same batch are identical, despite the data points themselves being distinct.