The masked_index is not working properly if dim != 0. One of my colleagues suggest this function:
def masked_index(input, dim, mask):
	assert len(mask.size())==1 and input.size(dim)==mask.size(0), \
		'{}!=1 or {}!={}'.format(len(mask.size()), input.size(dim), mask.size(0))
	indices = torch.arange(0,mask.size(0))[mask].long()
	return input.index_select(dim, indices)

using torch.arange, we can get the corresponding indices easily. (ref. range vs. arange 18)