vmapped the mmd loss function

train.py

@@ -563,7 +563,7 @@ def main():
 
     def regulariser_loss(latent_codes, rng):
         true_samples = jax.random.normal(rng, latent_codes.shape)
-        return compute_mmd(true_samples, latent_codes)
+        return jax.vmap(compute_mmd)(true_samples, latent_codes)
 
     def loss_fn(logits, labels, latent_codes, regulariser_rng):
         shift_logits = logits[..., :-1, :]