For my variable length input I have attempted to use a mask to avoid padding messing with the batch statistics.
x[~mask] = self.norm(x[~mask])
Here in this toy example x is of dimensionality (batch,embedding) and the mask is of dimensonality (batch) and is true where real data is and false where padding is. I have implemented this strategy and it seems extremely unstable. Is there something wrong with the way I’m going about this? Is there a better way to do a “partial” batch norm like this where padding is ignored?