added rng key to the eval_step

train.py

@@ -597,7 +597,7 @@ def main():
     def eval_step(params, batch):
         labels = batch.pop("labels")
         logits, latent_codes = model(**batch, params=params, train=False)[:2]
-        loss = loss_fn(logits, labels, latent_codes)
+        loss = loss_fn(logits, labels, latent_codes, rng)
 
         # summarize metrics
         metrics = {"loss": loss}