Repository: Lightning-AI/pytorch-lightning

Issue:
0.8.2 calls backward on '_GeneratorContextManager'
## 🐛 Bug

0.8.2 calls backward on '_GeneratorContextManager' and crashes training.
0.8.1 works correctly. my `training_step` returns `{'loss':loss, 'log':{'learn_rate':self.lr}}`

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 538, in ddp_train
    self.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1100, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 370, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 452, 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 630, in run_training_batch
    self.hiddens
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 804, in optimizer_closure
    model_ref.backward(self, closure_loss, optimizer, opt_idx)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/core/hooks.py", line 189, in backward
    loss.backward()
AttributeError: '_GeneratorContextManager' object has no attribute 'backward'
```

### Expected behavior
backward is called on the loss and training runs correctly


Discussion:
did you override optimizer step?
could you try master? we just pushed a fix to a typo we had
Can confirm this happens on 0.8.3
ok. Can you post a colab example that replicates this?
@Anjum48 @s-rog 
colab please
@williamFalcon my optimizer step was untouched, I can't run more testing atm but I'll get to it as soon as I can
@williamFalcon Hi I also encountered this, with normal Adam optimizer. I don't have a colab to replicate this atm but from what I saw earlier, this can be replicated with any setting as long as the Trainer is set to precision=16 when using Apex. Under this condition, the following lines from training_loop.py and hooks.py will run:

`if self.precision == 16 and not self.on_tpu
     closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx) `

`scaled_loss = amp.scale_loss(unscaled_loss, optimizer)`

will cause the closure_loss be a _GeneratorContextManager object. Which then cannot have a **backward()** method.

It seems under the current design, pytorch lighting's **scale_loss** function can only be used as a context?

@williamFalcon Here's a colab example (my first time using colab so let me know if you have issues seeing it) https://colab.research.google.com/drive/1G08jVDpx-T-5HE2c89RLJdq4u67mM2-o?usp=sharing

I suspect the issue lies with Apex AMP as suggested above by @aeryen
ummm. I think this is an apex issue. I can't replicate it with 16-bit native.

![image](https://user-images.githubusercontent.com/3640001/86135032-4c97ff80-bab8-11ea-942e-ffaae17aff07.png)

@aeryen min share a minimal example to reproduce?
hi sorry for the delay: https://colab.research.google.com/drive/1rjaRRwgBTm4CKPfe9po_WSxnKqY4jDRv?usp=sharing
I agree this is an apex issue, i.e. only occur when NATIVE_AMP_AVALAIBLE is false in the hooks.py

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
@@ -802,9 +802,22 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
             if self.precision == 16 and not self.on_tpu:
                 closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx)
 
+                # enter amp context
+                if not NATIVE_AMP_AVALAIBLE:
+                    context = closure_loss
+                    closure_loss = closure_loss.__enter__()
+
             # do backward pass
             model_ref.backward(self, closure_loss, optimizer, opt_idx)
 
+            # exit amp context
+            if self.precision == 16 and not NATIVE_AMP_AVALAIBLE:
+                a, b, c = None, None, None
+                error = context.__exit__(a, b, c)
+                if error:
+                    rank_zero_warn(a, b, c)
+                    raise Exception('apex unscale error')
+
             # once backward has been applied, release graph
             closure_loss = closure_loss.detach()
             training_step_output.batch_loss = training_step_output.batch_loss.detach()


Test patch: