Repository: coqui-ai/TTS

Issue:
multiband_melgan Vocoder Fails on Step 10000 With KeyError: 'avg_loss_0'
### Describe the bug

[multiband_melgan](https://github.com/coqui-ai/TTS/blob/dev/recipes/ljspeech/multiband_melgan/train_multiband_melgan.py) vocoder trained on [en](https://data.solak.de/data/Training/stt_tts/en_US.tgz), specifically, `"by_book/male/elliot_miller/pirates_of_ersatz/"` (tried another language, failed with the same error), fails on step 10000 with KeyError: 'avg_loss_0'.

### To Reproduce

Run train_multiband_melgan.py on en_US dataset.

### Expected behavior

No Errors.

### Logs

_No response_

### Environment

```shell
{
    "CUDA": {
        "GPU": [
            "Tesla T4",
            "Tesla T4",
            "Tesla T4",
            "Tesla T4"
        ],
        "available": true,
        "version": "10.2"
    },
    "Packages": {
        "PyTorch_debug": false,
        "PyTorch_version": "1.11.0+cu102",
        "TTS": "0.6.2",
        "numpy": "1.21.6"
    },
    "System": {
        "OS": "Linux",
        "architecture": [
            "64bit",
            "ELF"
        ],
        "processor": "x86_64",
        "python": "3.7.13",
        "version": "#1 SMP Tue Apr 26 20:14:22 UTC 2022"
    }
}
```


### Additional context

_No response_


Discussion:
Can you try the dev branch? I think we solve it there.
@erogol could you clarify which project dev branch has the fix - TTS or Trainer? 

It looks like the exception is coming from the Trainer, I am using TTS dev branch and still getting the same error.  Should I checkout and build Trainer locally?

Thank you!
🐸TTS dev 
Thank you.  I am on TTS-dev, still seeing this exception.  Full trace:

```
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/trainer/trainer.py", line 1489, in fit
    self._fit()
  File "/usr/local/lib/python3.7/dist-packages/trainer/trainer.py", line 1473, in _fit
    self.train_epoch()
  File "/usr/local/lib/python3.7/dist-packages/trainer/trainer.py", line 1252, in train_epoch
    _, _ = self.train_step(batch, batch_num_steps, cur_step, loader_start_time)
  File "/usr/local/lib/python3.7/dist-packages/trainer/trainer.py", line 1184, in train_step
    target_avg_loss = self._pick_target_avg_loss(self.keep_avg_train)
  File "/usr/local/lib/python3.7/dist-packages/trainer/trainer.py", line 1673, in _pick_target_avg_loss
    return keep_avg_target[f"avg_{self.config.target_loss}"]
  File "/usr/local/lib/python3.7/dist-packages/trainer/generic_utils.py", line 100, in __getitem__
    return self.avg_values[key]
KeyError: 'avg_loss_0'
```
Im having the same Issue in univenet :)
I'll check myself. 
Thank you
Thank you, @erogol !
@erogol Anything you found out? I cloned the repo (TTS-dev) again and tried it again with the original LJSpeech dataset and without modifications. The error unfortunately stays the same.
error is coming because of save_step
there are two issues
 one in trainer 
target_loss is not present in  keep_avg_target.avg_values 
```
    def _pick_target_avg_loss(self, keep_avg_target: KeepAverage) -> Dict:
        """Pick the target loss to compare models"""
        target_avg_loss = None

        # return if target loss defined in the model config
        if "target_loss" in self.config and self.config.target_loss:
            if f"avg_{self.config.target_loss}" in keep_avg_target.avg_values.keys():
                return keep_avg_target[f"avg_{self.config.target_loss}"]
            else:
                print(keep_avg_target.avg_values.keys())
                return keep_avg_target[f"avg_loss_1"]
```
another in gan.py _log 
output[0] is coming as None
```
def _log(name: str, ap: AudioProcessor, batch: Dict, outputs: Dict) -> Tuple[Dict, Dict]:
        """Logging shared by the training and evaluation.

        Args:
            name (str): Name of the run. `train` or `eval`,
            ap (AudioProcessor): Audio processor used in training.
            batch (Dict): Batch used in the last train/eval step.
            outputs (Dict): Model outputs from the last train/eval step.

        Returns:
            Tuple[Dict, Dict]: log figures and audio samples.
        """
        if outputs[0] == None:
            y_hat = outputs[1]["model_outputs"]
        else:
            y_hat = outputs[0]["model_outputs"]
        y = batch["waveform"]
        figures = plot_results(y_hat, y, ap, name)
        sample_voice = y_hat[0].squeeze(0).detach().cpu().numpy()
        audios = {f"{name}/audio": sample_voice}
        return figures, audios
```
i have put a workaround for my usecase 
@manmay-nakhashi Oh wow thank you. Is your posted code your workaround?
@LaughingC0ffin yes, can see which loss do you want when loss_0 is not present in the dictionary . 
haven't checked yet and I won't be able to at least for the next two weeks, unfortunately. 

If there is a fix shoot a PR 
@manmay-nakhashi thank you! 

I tested your fix and it is passing the 10K CHECKPOINT with a saving step.  

Resolving patch:
diff --git a/TTS/vocoder/models/gan.py b/TTS/vocoder/models/gan.py
--- a/TTS/vocoder/models/gan.py
+++ b/TTS/vocoder/models/gan.py
@@ -185,8 +185,7 @@ def train_step(self, batch: Dict, criterion: Dict, optimizer_idx: int) -> Tuple[
             outputs = {"model_outputs": self.y_hat_g}
         return outputs, loss_dict
 
-    @staticmethod
-    def _log(name: str, ap: AudioProcessor, batch: Dict, outputs: Dict) -> Tuple[Dict, Dict]:
+    def _log(self, name: str, ap: AudioProcessor, batch: Dict, outputs: Dict) -> Tuple[Dict, Dict]:
         """Logging shared by the training and evaluation.
 
         Args:
@@ -198,7 +197,7 @@ def _log(name: str, ap: AudioProcessor, batch: Dict, outputs: Dict) -> Tuple[Dic
         Returns:
             Tuple[Dict, Dict]: log figures and audio samples.
         """
-        y_hat = outputs[0]["model_outputs"]
+        y_hat = outputs[0]["model_outputs"] if self.train_disc else outputs[1]["model_outputs"]
         y = batch["waveform"]
         figures = plot_results(y_hat, y, ap, name)
         sample_voice = y_hat[0].squeeze(0).detach().cpu().numpy()


Test patch: