Repository: continualai/avalanche
Issue: EWC fails with RNN's using CUDA

🐛 **Describe the bug**

When training a model which contains an RNN/LSTM/GRU layer using the EWC strategy using CUDA, an error is raised since [RNN-like layers in PyTorch do not support](https://discuss.pytorch.org/t/when-should-we-set-torch-backends-cudnn-enabled-to-false-especially-for-lstm/106571/6) `backward` calls on CUDA devices when in `eval` mode.

#### Trace:

```
  File "/.../avalanche/training/strategies/base_strategy.py", line 248, in train
    self.train_exp(self.experience, eval_streams, **kwargs)
  File "/.../avalanche/training/strategies/base_strategy.py", line 304, in train_exp
    self.after_training_exp(**kwargs)
  File "/.../avalanche/training/strategies/base_strategy.py", line 536, in after_training_exp
    p.after_training_exp(self, **kwargs)
  File "/.../avalanche/training/plugins/ewc.py", line 92, in after_training_exp
    importances = self.compute_importances(strategy.model,
  File "/.../avalanche/training/plugins/ewc.py", line 123, in compute_importances
    loss.backward()
  File "/.../torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/.../torch/autograd/__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: cudnn RNN backward can only be called in training mode
```

#### Source of error:

https://avalanche-api.continualai.org/_modules/avalanche/training/plugins/ewc/#EWCPlugin

```python
    def compute_importances(self, model, criterion, optimizer,
                            dataset, device, batch_size):
        """
        Compute EWC importance matrix for each parameter
        """

        model.eval()                                       # <----------- Here eval mode is called

        # list of list
        importances = zerolike_params_dict(model)
        dataloader = DataLoader(dataset, batch_size=batch_size)
        for i, (x, y, task_labels) in enumerate(dataloader):
            x, y = x.to(device), y.to(device)

            optimizer.zero_grad()
            out = avalanche_forward(model, x, task_labels)
            loss = criterion(out, y)
            loss.backward()                                 # <----------- Given above, here causes the error
```

🐜 **To Reproduce**
- Use a machine with a CUDA capable GPU.
- Create a model containing an RNN-like layer (RNN/LSTM/GRU).
- Wrap it in the EWC strategy.
- Set the `device` to the GPU.
- Train the model (on any dataset).

🐝 **Expected behavior**
RNN models to run with EWC without error on GPU.

🦋 **Additional context**

This appears to be the opposite behaviour of the `EWCPlugin` as defined in:

https://avalanche-api.continualai.org/_modules/avalanche/training/plugins/#EWCPlugin

```python
    def compute_importances(self, model, criterion, optimizer,
                            dataset, device, batch_size):
        """
        Compute EWC importance matrix for each parameter
        """

        model.train()   
```

🐞 **Potential fix**

If this is just a typo, the above would be fixed by changing `model.eval()` to `model.train()`.

Changed files:
avalanche/training/plugins/ewc.py

diff --git a/avalanche/training/plugins/ewc.py b/avalanche/training/plugins/ewc.py
index 20b7eb76..29de1c99 100644
--- a/avalanche/training/plugins/ewc.py
+++ b/avalanche/training/plugins/ewc.py
@@ -1,5 +1,6 @@
 from collections import defaultdict
 from typing import Dict, Tuple
+import warnings
 
 import torch
 from torch import Tensor
@@ -112,6 +113,19 @@ class EWCPlugin(StrategyPlugin):
 
         model.eval()
 
+        # Set RNN-like modules on GPU to training mode to avoid CUDA error
+        if device == 'cuda':
+            for module in model.modules():
+                if isinstance(module, torch.nn.RNNBase):
+                    warnings.warn(
+                        'RNN-like modules do not support '
+                        'backward calls while in `eval` mode on CUDA '
+                        'devices. Setting all `RNNBase` modules to '
+                        '`train` mode. May produce inconsistent '
+                        'output if such modules have `dropout` > 0.'
+                    )
+                    module.train()
+
         # list of list
         importances = zerolike_params_dict(model)
         dataloader = DataLoader(dataset, batch_size=batch_size)