show where to add reg loss

train.py

@@ -56,17 +56,15 @@ from flax.jax_utils import unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
 from transformers import (
-    CONFIG_MAPPING,
     FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
-    AutoConfig,
     AutoTokenizer,
-    FlaxAutoModelForCausalLM,
     HfArgumentParser,
     TrainingArguments,
     is_tensorboard_available,
 )
 from transformers.testing_utils import CaptureLogger
 
+from model.t5_vae import Funnel_T5_VAE_Model
 from model.config import T5_VAE_Config
 
 
@@ -526,10 +524,11 @@ def main():
     # Setup train state
     state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
 
-    def loss_fn(logits, labels):
+    def loss_fn(logits, labels, latent_codes):
         shift_logits = logits[..., :-1, :]
         shift_labels = labels[..., 1:]
         loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
+        # TODO add reg loss here
         return loss.mean()
 
     # Define gradient update step fn
@@ -538,8 +537,8 @@ def main():
 
         def compute_loss(params):
             labels = batch.pop("labels")
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-            loss = loss_fn(logits, labels)
+            logits, latent_codes = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[:2]
+            loss = loss_fn(logits, labels, latent_codes)
             return loss
 
         grad_fn = jax.value_and_grad(compute_loss)
@@ -556,8 +555,8 @@ def main():
     # Define eval fn
     def eval_step(params, batch):
         labels = batch.pop("labels")
-        logits = model(**batch, params=params, train=False)[0]
-        loss = loss_fn(logits, labels)
+        logits, latent_codes = model(**batch, params=params, train=False)[:2]
+        loss = loss_fn(logits, labels, latent_codes)
 
         # summarize metrics
         metrics = {"loss": loss}