[terminal_command] TERMINAL ()
             psnr=psnr,             ssim=ssim,+            loss=mse,         ) -        loss = mse-        metrics["loss"] = loss--        return loss, (outputs["recon"], metrics)+        return mse, (outputs["recon"], metrics)      @nnx.jit(donate_argnums=0)     def train_step(diff --git a/jasmine/train_tokenizer_vqvae.py b/jasmine/train_tokenizer_vqvae.pyindex c76396d..0d62275 100644--- a/jasmine/train_tokenizer_vqvae.py+++ b/jasmine/train_tokenizer_vqvae.py@@ -343,6 +343,15 @@ def main(args: Args) -> None:         outputs = model(inputs, training=training)         outputs["recon"] = outputs["recon"].astype(jnp.float32)         mse = jnp.square(gt - outputs["recon"]).mean()+        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()+        commitment_loss = jnp.square(+            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])+        ).mean()+        loss = mse + q_loss + args.vq_beta * commitment_loss+        _, index_counts = jnp.unique_counts(+            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0+        )+        codebook_usage = (index_counts != 0).mean()          gt_clipped = gt.clip(0, 1).reshape(-1, *gt.shape[2:])         recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])@@ -353,22 +362,12 @@ def main(args: Args) -> None:             mse=mse,             psnr=psnr,             ssim=ssim,+            loss=loss,+            q_loss=q_loss,+            commitment_loss=commitment_loss,+            codebook_usage=codebook_usage,         ) -        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()-        commitment_loss = jnp.square(-            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])-        ).mean()-        loss = mse + q_loss + args.vq_beta * commitment_loss-        _, index_counts = jnp.unique_counts(-            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0-        )-        codebook_usage = (index_counts != 0).mean()-        metrics["loss"] = loss-        metrics["q_loss"] = q_loss-        metrics["commitment_loss"] = commitment_loss-        metrics["codebook_usage"] = codebook_usage-         return loss, (outputs["recon"], metrics)      @nnx.jit(donate_argnums=0)@@ -405,15 +404,13 @@ def main(args: Args) -> None:         (loss, (recon, metrics)) = tokenizer_loss_fn(tokenizer, inputs, training=False)         return loss, recon, metrics -    def calculate_validation_metrics(val_dataloader, tokenizer, rng):+    def calculate_validation_metrics(val_dataloader, tokenizer):         step = 0         loss_per_step = []         metrics_per_step = []         batch = None         recon = None         for batch in val_dataloader:-            rng, _rng_mask = jax.random.split(rng, 2)-            batch["rng"] = _rng_mask             loss, recon, metrics = val_step(tokenizer, batch)             loss_per_step.append(loss)             metrics_per_step.append(metrics)@@ -454,9 +451,7 @@ def main(args: Args) -> None:             for elem in val_iterator         )     if jax.process_index() == 0:-        rng, _rng = jax.random.split(rng)         first_batch = next(dataloader_train)-        first_batch["rng"] = _rng         compiled = train_step.lower(optimizer, first_batch).compile()         print_compiled_memory_stats(compiled.memory_analysis())         print_compiled_cost_analysis(compiled.cost_analysis())@@ -467,8 +462,6 @@ def main(args: Args) -> None:     while step < args.num_steps:         for batch in dataloader_train:             # --- Train step ----            rng, _rng = jax.random.split(rng)-            batch["rng"] = _rng             loss, recon, metrics = train_step(optimizer, batch)             if step == first_step:                 print_mem_stats("After params initialized")@@ -478,9 +471,8 @@ def main(args: Args) -> None:             val_results = {}             if dataloader_val and step % args.val_interval == 0:                 print("Calculating validation metrics...")-                rng, _rng_mask_val = jax.random.split(rng, 2)                 val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(-                    dataloader_val, optimizer.model, _rng_mask_val+                    dataloader_val, optimizer.model                 )                 print(f"Step {step}, validation loss: {val_metrics['val_loss']}")                 val_results = {(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ git commit -am "codecleanup"black....................................................................Passed[chore/refactor-diffusion-branch fac06ff] codecleanup 3 files changed, 18 insertions(+), 28 deletions(-)(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ git pushEnumerating objects: 13, done.Counting objects: 100% (13/13), done.Delta compression using up to 224 threadsCompressing objects: 100% (7/7), done.Writing objects: 100% (7/7), 713 bytes | 356.00 KiB/s, done.Total 7 (delta 6), reused 0 (delta 0), pack-reused 0remote: Resolving deltas: 100% (6/6), completed with 6 local objects.To github.com:p-doom/jasmine.git   75e6171..fac06ff  chore/refactor-diffusion-branch -> chore/refactor-diffusion-branch(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ git branch  chore/cleanup-diffusion-branch* chore/refactor-diffusion-branch  feat/diffusion-backend  feat/diffusion-forcing  feat/diffusion-jittable-sampling  feat/diffusion-lookup-embed  feat/diffusion-prepend-conditioning  feat/diffusion-ramp-weight  feat/diffusion-train-from-mae  feat/diffusion-x-loss  feat/masked-autoencoder  feat/notebook-detect-arrayrecord-duplicates  feat/shortcut-forcing  feat/sliding-window-sampling  main  prepend-action-maskgit  vizdoom-dataset(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ queue(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ fqueue(jasmine) [mihir.mahajan@hai003.haicore.berlin:~/Projects/jasmine] $ salloc: Job 32900 has exceeded its time limit and its allocation has been revoked.  srun: Job step aborted: Waiting up to 32 seconds for job step to finish.[2025-10-27T19:19:28.011] error: *** STEP 32900.interactive ON hai003 CANCELLED AT 2025-10-27T19:19:28 DUE TO TIME LIMIT ***srun: error: hai003: task 0: Killed