Mask `z` axis in reconstruction loss

The `z` axis in the input is quite noisy, and maybe not worth to learn.
Keeping it in the input would always be good, since the model can decide to learn from it or not, but possibly removing it from the loss (or weighting it down) can allow the model to better reconstruct the `x`/`y` axes.
