Indexing with boolean masks and ellipsis leads to inconsistent behavior.
MWE:
import torch
t = torch.arange(12).view(3,2,-1)
mask = torch.BoolTensor([[True, False], [False, True], [False, False]])

the expression
t[mask][...,0]

returns
>>> tensor([0, 6])

whereas
t[mask,..., 0]

returns
>>> tensor([0, 5])

With numpy arrays, both [mask,..., 0] and [mask][..., 0] return >>> array([0, 6]).