Some little remarks:

you could use mask.dim() rather than len(mask.size()) I think 

you could use mask.nonzero() rather than torch.arange(0,mask.size(0))[mask].long(), it gives the same result but is way more explicit !

All in all, I’m not sure this require a separate function, since you can simply write:
input.index_select(dim, mask.nonzero())

Hope this helps !