FIX Prefix tuning test w/ rotary emb on multi GPU (#2311)

See
huggingface/transformers#35235 (comment)
for context.

There has been a refactor in transformers that resulted in the rotary
embedding of Mistral (and probably others) moving to the model level.
This led to a device map used in one of the tests to being incorrect.
This PR fixes the device map.

Note that this fix doesn't really have anything to do with prefix
tuning, the error occurred even before prefix tuning is used.

tests/test_gpu_examples.py

@@ -4043,6 +4043,7 @@ def test_prefix_tuning_multiple_devices_decoder_model(self):
             "model.layers.0": 0,
             "model.layers.1": 1,
             "model.norm": 1,
+            "model.rotary_emb": 1,
             "lm_head": 1,
         }
         model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map)