Repository: Lightning-AI/pytorch-lightning

Issue:
Trainer crashed when optimizer frequency is defined. 
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior:

Run the following code:
https://gist.github.com/24hours/ec67de5384bb05e28544d580ae424639

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```
Traceback (most recent call last):
  File "pl_bug.py", line 40, in <module>
    trainer.fit(mnist_model, train_loader) 
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/states.py", line 48, in wrapped_fn
    result = fn(self, *args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1073, in fit
    results = self.accelerator_backend.train(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/accelerators/gpu_backend.py", line 51, in train
    results = self.trainer.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1239, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 394, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 491, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 883, in run_training_batch
    batch_outputs[opt_idx].append(opt_closure_result.training_step_output_for_epoch_end)
```

#### Code sample

```python
    def configure_optimizers(self):
        optimizer_G = torch.optim.Adam(self.parameters(), lr=0.1, weight_decay=1e-5)
        optimizer_D = torch.optim.Adam(self.parameters(), lr=0.1, weight_decay=1e-5)

        return [
                {'optimizer': optimizer_D, 'frequency': 5},
                {'optimizer': optimizer_G, 'frequency': 1}
            ]
```

the culprit is 'frequency' : 5, removing the line will allow trainer to run smoothly. 

https://pytorch-lightning.readthedocs.io/en/latest/lightning-module.html?highlight=optimizers#configure-optimizers
The definition is correct according to this documentation. 

### Expected behavior

Model should train without crash. 
The code work in 0.8.5 environment. 

### Environment

* CUDA:
	- GPU:
		- GeForce GTX TITAN X
	- available:         True
	- version:           11.0
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0a0+9907a3e
	- pytorch-lightning: 0.9.0
	- tensorboard:       2.2.0
	- tqdm:              4.48.2
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- 
	- processor:         x86_64
	- python:            3.6.10
	- version:           #110-Ubuntu SMP Tue Jun 23 02:39:32 UTC 2020

  


Discussion:
Hi! thanks for your contribution!, great first issue!
this is the case with multiple optimizers, you need to spec them.. so you would prefer having default 1 if freq is not specified?
the exception occur because trainer_loop.py incorrect count number of optimizer if `frequency` is defined. Since this configuration work in version 0.8.5, this look like regression error. 

Unless of course if the configuration is unsupported in version 0.9.0  

> so you would prefer having default 1 if freq is not specified?

That is a totally different case I think. If the frequency is not specified we run `train_step` for both optimizers but if it is specified to 1 for both then in such case it will run 1st batch for opt_1, 2nd for opt_2, 3rd for opt_1, 4th for opt_2...

A simple fix here can be:
https://github.com/PyTorchLightning/pytorch-lightning/blob/a7705c8677b9e2b5105e26a38db9d1b650182576/pytorch_lightning/trainer/training_loop.py#L872

```python
if len(batch_outputs) == 1:  # when frequencies are defined
    batch_outputs[0].append(opt_closure_result.training_step_output_for_epoch_end)
else:  # no frequencies
    batch_outputs[opt_idx].append(opt_closure_result.training_step_output_for_epoch_end)
```
ok, can someone write a test and submit a PR for this? show the test failing on master first.
@awaelchli  or @rohitgr7 ?

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -476,7 +476,8 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
                 self.accumulated_loss.append(opt_closure_result.loss)
 
                 # track all the outputs across all steps
-                batch_outputs[opt_idx].append(opt_closure_result.training_step_output_for_epoch_end)
+                batch_opt_idx = opt_idx if len(batch_outputs) > 1 else 0
+                batch_outputs[batch_opt_idx].append(opt_closure_result.training_step_output_for_epoch_end)
 
                 # ------------------------------
                 # BACKWARD PASS


Test patch:
diff --git a/tests/base/model_optimizers.py b/tests/base/model_optimizers.py
--- a/tests/base/model_optimizers.py
+++ b/tests/base/model_optimizers.py
@@ -34,6 +34,14 @@ def configure_optimizers__multiple_optimizers(self):
         optimizer2 = optim.Adam(self.parameters(), lr=self.learning_rate)
         return optimizer1, optimizer2
 
+    def configure_optimizers__multiple_optimizers_frequency(self):
+        optimizer1 = optim.Adam(self.parameters(), lr=self.learning_rate)
+        optimizer2 = optim.Adam(self.parameters(), lr=self.learning_rate)
+        return [
+            {'optimizer': optimizer1, 'frequency': 1},
+            {'optimizer': optimizer2, 'frequency': 5}
+        ]
+
     def configure_optimizers__single_scheduler(self):
         optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
         lr_scheduler = optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.1)
diff --git a/tests/trainer/test_optimizers.py b/tests/trainer/test_optimizers.py
--- a/tests/trainer/test_optimizers.py
+++ b/tests/trainer/test_optimizers.py
@@ -240,3 +240,18 @@ def configure_optimizers(self):
     )
     result = trainer.fit(model)
     assert result == 1
+
+
+def test_configure_optimizers_with_frequency(tmpdir):
+    """
+    Test that multiple optimizers work when corresponding frequency is set.
+    """
+    model = EvalModelTemplate()
+    model.configure_optimizers = model.configure_optimizers__multiple_optimizers_frequency
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1
+    )
+    result = trainer.fit(model)
+    assert result