Fix transformer_moe model has wrong logic in pre/postprocessing (#1233)

Author: twilightdema

tensor2tensor/models/research/transformer_moe.py

@@ -93,8 +93,8 @@ def prepostprocess(fct):
       """Apply processing and capture the extra loss."""
       @expert_utils.add_var_scope()
       def decorated(x, *args, **kwargs):
-        x = dp_preprocess(x)
-        y, loss = fct(x, *args, **kwargs)
+        x_preprocessed = dp_preprocess(x)
+        y, loss = fct(x_preprocessed, *args, **kwargs)
         cache["extra_loss"] += loss
         return dp_postprocess(x, y)
       return decorated