gpt2 model parallelization test failed
### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: Linux-5.13.0-1015-gcp-x86_64-with-glibc2.17
- Python version: 3.8.12
- Huggingface_hub version: 0.4.0
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes
```


### Who can help?

In main branch, ``` pytest tests/models/gpt2/test_modeling_gpt2.py ``` failed at ```test_model_parallelization```.

```
 def test_model_parallelization(self):

            ...

            # Assert that the memory use on all devices is higher than it was when loaded only on CPU
            for n in range(torch.cuda.device_count()):
>               self.assertGreater(memory_after_parallelization[n], memory_at_start[n])
E               AssertionError: 0 not greater than 0

tests/test_modeling_common.py:2069: AssertionError
```

I'm implementing model parallelization on OPT, but there is the same problem. (https://github.com/huggingface/transformers/pull/17245)

But, it works on trainer.

@patrickvonplaten

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

``` pytest tests/models/gpt2/test_modeling_gpt2.py ```

### Expected behavior

```shell
memory_at_start : [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
memory_after_parallelization : [179, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 0, 0, 0, 0]

Number of my gpu devices is 16. but len(self.h) at gpt2 is 12.
```
