Repository: huggingface/transformers
Pull request: fix(models): Fix suno/bark-small CPU offload device mismatch causing CI failures

### What does this PR do?

→ `BarkCausalModel` and `BarkFineModel` incorrectly inferred the device, causing device mismatches when `enable_cpu_offload()` is used. For more details on reproducing the bug, please visit the linked issue!

Fixes #43606.

### Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you fix any necessary existing tests?

**CI Failure:**

<img width="1896" height="174" alt="Image" src="https://github.com/user-attachments/assets/d06f6d23-0d52-4ed0-ba5a-98f78f92a0d0" />

**Current Output:**

<img width="1791" height="1084" alt="Image" src="https://github.com/user-attachments/assets/5a84a66c-d134-4fa3-ae0a-57b746310ce3" />

**Output After the Fix:**

<img width="500" height="500" alt="Image" src="https://github.com/user-attachments/assets/90548ad8-8af7-4331-af9e-33bc4c8a4f8a" />

Changed files:
filename: src/transformers/models/bark/modeling_bark.py
patch: @@ -474,8 +474,6 @@ def forward(
         batch_size = input_embeds.shape[0]
         seq_length = input_shape[-1]
 
-        device = input_ids.device if input_ids is not None else input_embeds.device
-
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(
@@ -487,11 +485,18 @@ def forward(
             past_key_values = DynamicCache(config=self.config)
 
         past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        input_embeds = input_embeds.to(self.position_embeds_layer.weight.device)
 
         if position_ids is None:
-            position_ids = torch.arange(past_length, seq_length + past_length, dtype=torch.long, device=device)
+            position_ids = torch.arange(
+                past_length,
+                seq_length + past_length,
+                dtype=torch.long,
+                device=self.position_embeds_layer.weight.device,
+            )
             position_ids = position_ids.unsqueeze(0)  # shape (1, seq_length)
 
+        position_ids = position_ids.to(self.position_embeds_layer.weight.device)
         position_embeds = self.position_embeds_layer(position_ids)  # position embeddings of shape (1, t, n_embd)
 
         # Attention mask.
@@ -505,6 +510,7 @@ def forward(
                 # [bsz, to_seq_length] -> [bsz, 1, 1, to_seq_length]
                 # from_seq_length is 1 to easily broadcast
                 attention_mask = _prepare_4d_attention_mask(attention_mask, input_embeds.dtype, tgt_len=1)
+                attention_mask = attention_mask.to(input_embeds.device)
 
         hidden_states = self.drop(input_embeds + position_embeds)
         output_shape = input_shape + (hidden_states.size(-1),)
@@ -1084,12 +1090,15 @@ def forward(
         batch_size = input_embeds.shape[0]
         seq_length = input_shape[1]
 
-        device = input_ids.device if input_ids is not None else input_embeds.device
+        input_embeds = input_embeds.to(self.position_embeds_layer.weight.device)
 
         if position_ids is None:
-            position_ids = torch.arange(0, seq_length, dtype=torch.long, device=device)
+            position_ids = torch.arange(
+                0, seq_length, dtype=torch.long, device=self.position_embeds_layer.weight.device
+            )
             position_ids = position_ids.unsqueeze(0)  # shape (1, seq_length)
 
+        position_ids = position_ids.to(self.position_embeds_layer.weight.device)
         position_embeds = self.position_embeds_layer(position_ids)  # position embeddings of shape (1, t, n_embd)
 
         # Attention mask.
@@ -1102,6 +1111,7 @@ def forward(
                 # [bsz, to_seq_length] -> [bsz, 1, 1, to_seq_length]
                 # from_seq_length is 1 to easily broadcast
                 attention_mask = _prepare_4d_attention_mask(attention_mask, input_embeds.dtype, tgt_len=1)
+                attention_mask = attention_mask.to(input_embeds.device)
 
         hidden_states = self.drop(input_embeds + position_embeds)
         output_shape = input_shape + (hidden_states.size(-1),)