🐛 Support frequencies for multiple optimizers

see Lightning-AI#3229
rephrase-ai · Aug 29, 2020 · da85f45 · da85f45
1 parent 941baaa
commit da85f45
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
@@ -872,7 +872,8 @@ def run_training_batch(self, batch, batch_idx):
                 self.batch_loss_value.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_output_idx = 0 if len(batch_outputs) == 1 else opt_idx
+                batch_outputs[_batch_output_idx].append(opt_closure_result.training_step_output_for_epoch_end)
 
                 # ------------------------------
                 # BACKWARD PASS