Fix device issue in `SwitchTransformers`
# What does this PR do?

Need a tiny fix after #24300.

Currently, we have a failure


```bash

self = <tests.models.switch_transformers.test_modeling_switch_transformers.SwitchTransformersEncoderOnlyModelTest testMethod=test_multi_gpu_data_parallel_forward>

    @staticmethod
    def forward(ctx, target_device, dim, *inputs):
>       assert all(i.device.type != 'cpu' for i in inputs), (
            'Gather function not implemented for CPU tensors'
        )
E       AssertionError: Gather function not implemented for CPU tensors

/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/_functions.py:56: AssertionError
```
_The documentation is not available anymore as the PR was closed or merged._<|||||>Merge now. Don't hesitate to leave comments if any @ArthurZucker .