Tests failing on GPU

**Reproducing**

On a machine with a GPU:

```
$ CUDA_VISIBLE_DEVICES=-1 python -m pytest deepinv/tests/test_trainer.py
...
45 passed, 8 skipped, 42 warnings in 4.78s
```

```
$ python -m pytest deepinv/tests/test_trainer.py
...
24 failed, 21 passed, 8 skipped, 25 warnings in 4.13s
```