apply z-loss and aux-loss only during training

[xffxff]

After profiling, I discovered that the routing logic consumes a significant amount of time, primarily due to apply_z_loss and apply_aux_loss, which are unnecessary during inference´. Thus, we can eliminate these operations in the inference phase