Well, the point is that data and mask have shape n*8*2, and this is the reasonable shape for mu in model according to its construction process (gain_* also have the shape n*8*2). And data[~mask] (as well as mu[~mask]) is a single-dimensional vector with size equals to zeros count in mask, which makes it impossible to minibatch data[~mask] together with e.g. ix2 (which is n-sized).