Repository: wandb/wandb

Issue:
[CLI] RuntimeError: OrderedDict mutated during iteration
I am getting a `RuntimeError: OrderedDict mutated during iteration` error whenever I try to do _any forward_ calls with a model that I have previously called `wandb.watch` on. **This does not happen when I do not call `wandb.watch` on the model**. I have manually tested all wandb versions from `0.10.27` and upwards and can confirm that this error only starts to appear in `0.10.30` and is not fixed in `0.10.31`.

In this specific traceback I am also using PytorchLightning, however I was getting the same error in an other project without Lightning. I have been forced to downgrade `wandb` to `0.10.29` due to this.

Traceback:
```sh
Traceback (most recent call last):
  File "src/gan/train.py", line 83, in <module>
    trainer.fit(model, dm)
  File "/Users/user/prog/mp/genart/venv/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 458, in fit
    self._run(model)
 <... a bunch of irrelevant PytorchLightning stuff ...>
  File "/Users/user/prog/mp/genart/venv/lib/python3.7/site-packages/pytorch_lightning/core/memory.py", line 242, in summarize
    self._forward_example_input()
  File "/Users/user/prog/mp/genart/venv/lib/python3.7/site-packages/pytorch_lightning/core/memory.py", line 267, in _forward_example_input
    model(input_)
  File "/Users/user/prog/mp/genart/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/user/prog/mp/genart/src/gan/outer_gan.py", line 127, in forward
    generated_images = self.generator(z, labels)
  File "/Users/user/prog/mp/genart/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 892, in _call_impl
    self._forward_hooks.values()):
RuntimeError: OrderedDict mutated during iteration
```



Discussion:
Hi @konstantinjdobler thanks for reporting this. @borisdayma could you please take a look.
I'm seeing this as well using our W&B training callback in AllenNLP with `watch_model=True`.

![image](https://user-images.githubusercontent.com/8812459/121072347-f3375680-c785-11eb-8449-ff4bfdd72fad.png)
We had a same issue when running in 0.10.30. We had to revert back to an earlier version.
Just letting you know it's being worked on. Hopefully we'll have a fix for next week.

Resolving patch:
diff --git a/wandb/wandb_torch.py b/wandb/wandb_torch.py
--- a/wandb/wandb_torch.py
+++ b/wandb/wandb_torch.py
@@ -312,6 +312,7 @@ def __init__(self):
         # when we detect fastai v1 has been imported :(
         self._should_use_full_hooks = True
         self._graph_hooks = {}
+        self._graph_hooks_processed = []
         if "fastai" in sys.modules:
             fastai = util.get_module("fastai")
             if fastai.__version__.startswith("1."):
@@ -362,10 +363,12 @@ def after_forward_hook(module, input, output):
             # log graph and remove hook
             hook = self._graph_hooks.pop(id(module), None)
             if hook is not None:
-                hook.remove()
+                self._graph_hooks_processed.append(hook)
 
             if not self._graph_hooks:
                 # we went through the entire graph
+                for hook_processed in self._graph_hooks_processed:
+                    hook_processed.remove()
                 wandb.run.summary["graph_%i" % graph_idx] = self
 
         return after_forward_hook


Test patch: