fix device issue
Author: grahamannett
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

If using GPU's 48GB or below where model must be split across devices, the tensors are on different devices at this point. Issue is when using device_map="auto" as seen in examples: https://huggingface.co/adept/fuyu-8b/discussions/44#6544c5e6ee7bbb5952bdebfb

You get this issue
```
    output_embeddings[batch_idx, dst_indices] = continuous_embeddings[batch_idx][src_indices]
                                                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor
```

This seems to fix when running on 48GB cards but think there may still be other issues that only noticeable on 24GB.

## Who can review?

Probably @ArthurZucker as I believe I saw them working on this before.  


Comments:
this might need to be split into 3 lines as I edited the file online so no format/linter

Hey! I'll ping @molbap as he recently re-worked that part! Thanks for reporting


Hi @grahamannett, thanks! indeed this would cause errors, ideally it would be good to ensure all tensors are on the same device for that step. Before that step, the output of `FuyuProcessor` does send inputs to devices correctly, but at this stage the indexing can break it.
Can you try, before the actual batch loop, something like this:

```python
    device = continuous_embeddings[0].device
    word_embeddings = word_embeddings.to(device)
    image_patch_input_indices = image_patch_input_indices.to(device)
```
This is assuming `continuous_embeddings` is the largest tensor.


@molbap nah that does not work for me and continuous embeddings will most likely not be the largest tensor (especially if training), e.g.

```python
(Pdb) word_embeddings.shape
torch.Size([1, 1130, 4096])
(Pdb) image_patch_input_indices.shape
torch.Size([1, 1140])
(Pdb) continuous_embeddings[0].shape
torch.Size([792, 4096])
```

Believe there is also an issue with src_indices and dst_indices shapes if you are using multiple variable length inputs.  Not really related to this issue but if anyone runs into it the easiest way is to change the raise ValueError to something like
```
if src_indices.shape[0] > continuous_embeddings[batch_idx].shape[0]:
    src_indices = src_indices[: continuous_embeddings[batch_idx].shape[0]]
    dst_indices = dst_indices[: src_indices.shape[0]]
```

It probably is not necessary if the gather_continuous_embeddings is rewritten and would probably be quicker overall if it was (get rid of the for loop and unnecessary tensor allocations), but also given how large this model is and my resources it seems like very little of the actual time spent in forward/backward is related to this.

@ArthurZucker @molbap it might be worth refactoring a fair amount of the fuyu related stuff.  Ive ended up having to monkey patch various issues on both the model and processor.  

I have a list of what I've done to the model and processor but not sure if some of them are me just not understanding.  Latest one is if you feed a bbox into the processor where the value of one of the ints isnt in the vocab, for instance make a bbox at  1007, 1008, 1009, 1010 (i think the vocab ends at 999 but images are 1080x1200 or something?).  

For the bbox related issue, what happens specifically in your last example, raises a `KeyError`? I assume it's related to the original bounding box conversion function in here https://github.com/huggingface/transformers/blob/b54993aa94f744598129db32d7e1ec4e7c299099/src/transformers/models/fuyu/processing_fuyu.py#L175C1-L208C1, any ideas @pcuenca since I think you worked on that part?

 A note @grahamannett  if you have found other issues, feel free to open a PR/Issue with a description + code reproduction separately for each, would be helpful!

@grahamannett @molbap re: the bounding box problem, that's a good question! I had the same doubt when I ported the code, but I did not see any special code to handle those cases in the original repo, so I assumed all tokens had been added to the vocabulary. @grahamannett do you have a short snippet to reproduce?

@pcuenca can you point me to where you even ported the code from?  I don't really see anything similar in https://github.com/persimmon-ai-labs/adept-inference and fine-tuning with the current processor my results aren't good at all so would like to see what the original looked like.

Hi @grahamannett! Your original PR is laser-focused on a very specific problem (thanks for that! 🙌). To avoid derailing it, I'd suggest you open a new issue where we can discuss about the bounding box processing method and your observation that not all possible coordinates appear to be backed by single tokens.

I'll let @molbap work with you in assessing the best way to solve the problem your PR is trying to address.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.