Fix codec e2e tests (#15607)

Signed-off-by: Ryan <rlangman@nvidia.com>

Author: rlangman

tests/e2e_nightly/test_model_support_mel_codec_22khz_fullband_medium.py

@@ -69,7 +69,7 @@ def test_model_training_step():
         "audio_lens": torch.tensor([num_samples], device=d),
     }
 
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     generator_losses = []
     loss_mel_l1, loss_mel_l2 = model.mel_loss_fn(

tests/e2e_nightly/test_model_support_nvidia__audio_codec_22khz.py

@@ -68,7 +68,7 @@ def test_model_training_step():
         "audio_lens": torch.tensor([num_samples], device=d),
     }
 
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     generator_losses = []
     loss_mel_l1, loss_mel_l2 = model.mel_loss_fn(

tests/e2e_nightly/test_model_support_nvidia__audio_codec_44khz.py

@@ -73,7 +73,7 @@ def test_model_training_step():
     }
 
     # Forward pass through encoder -> VQ -> decoder (same as training_step).
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     generator_losses = []
 

tests/e2e_nightly/test_model_support_nvidia__low_frame_rate_speech_codec_22khz.py

@@ -76,7 +76,7 @@ def test_model_training_step():
     }
 
     # Run the shared forward pass used by training_step.
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     # Compute the mel reconstruction loss (l1 + l2) — same calls as training_step.
     loss_mel_l1, loss_mel_l2 = model.mel_loss_fn(

tests/e2e_nightly/test_model_support_nvidia__mel_codec_22khz.py

@@ -71,7 +71,7 @@ def test_model_training_step():
     }
 
     # _process_batch runs encoder -> (optional) vector quantizer -> decoder.
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     # Compute the L1 mel loss, which is always enabled (mel_loss_l1_scale > 0
     # by default) and is the primary reconstruction objective.

tests/e2e_nightly/test_model_support_nvidia__mel_codec_44khz.py

@@ -72,7 +72,7 @@ def test_model_training_step():
     }
 
     # _process_batch runs encoder -> (optional) vector quantizer -> decoder.
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     # Compute the L1 mel loss, which is always enabled (mel_loss_l1_scale > 0
     # by default) and is the primary reconstruction objective.

tests/e2e_nightly/test_model_support_nvidia__nemo_nano_codec_22khz_0_6kbps_12_5fps.py

@@ -66,7 +66,7 @@ def test_model_training_step():
     audio_len = torch.tensor([num_samples], dtype=torch.long, device=d)
     batch = {"audio": audio, "audio_lens": audio_len}
 
-    audio_ref, audio_ref_len, audio_gen, commit_loss, _ = model._process_batch(batch)
+    audio_ref, audio_ref_len, audio_gen, commit_loss, *_ = model._process_batch(batch)
 
     loss_mel_l1, loss_mel_l2 = model.mel_loss_fn(
         audio_real=audio_ref.float(),

tests/e2e_nightly/test_model_support_nvidia__nemo_nano_codec_22khz_1_78kbps_12_5fps.py

@@ -72,7 +72,7 @@ def test_model_training_step():
     # _process_batch runs encoder → (optional) quantizer → decoder and returns
     # the (padded) reference audio, its lengths, the generated audio, the
     # commit loss from the quantizer, and the encoded representation.
-    audio_ref, audio_ref_len, audio_gen, commit_loss, _ = model._process_batch(batch)
+    audio_ref, audio_ref_len, audio_gen, commit_loss, *_ = model._process_batch(batch)
 
     # Compute the mel reconstruction losses (the primary generator losses).
     # mel_loss_fn returns (loss_l1, loss_l2); both are scalar tensors.

tests/e2e_nightly/test_model_support_nvidia__nemo_nano_codec_22khz_1_89kbps_21_5fps.py

@@ -68,7 +68,7 @@ def test_model_training_step():
         "audio_lens": torch.tensor([num_samples], dtype=torch.long, device=d),
     }
 
-    audio, audio_len, audio_gen, commit_loss, codes = model._process_batch(batch)
+    audio, audio_len, audio_gen, commit_loss, codes, *_ = model._process_batch(batch)
 
     # Only compute mel reconstruction losses (skip discriminator).
     loss_mel_l1, loss_mel_l2 = model.mel_loss_fn(