[Tests] Fix ViTMAE integration test
# What does this PR do?

This PR fixes the integration test of ViTMAE, which only passed on CPU.

ViTMAE uses randomness inside (it creates a random boolean mask to indicate which patches to mask). I used `torch.manual_seed(2)` to make this deterministic, however I [learned](https://discuss.pytorch.org/t/random-seed-that-spans-across-devices/19735) that using the same seed across CPU/GPU will not result in the same random numbers (as both use a different random number generator).

Hence, this PR sets the expected slice depending on the device.

Huge thanks to @ydshieh for helping me out finding this.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15949). All of your documentation changes will be reflected on that endpoint.