Create the return value on device to avoid unnecessary copying from CPU

[mksit]

What does this PR do?

router_tuple = (torch.tensor([0], device=hidden_states.device),) introduces an unnecessary data copy from the CPU. I have changed it to create the return tensor on the device to avoid potential performance issues.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

[ArthurZucker]

Are you implying that torch.tensor([0], device=hidden_states.device) does not create the tensor on device=hidden_states.device ? The [0] is copied but it does not have a place from where it's copied no?

Do you have a reference for this the doc does not seem to indicate this:

device (torch.device, optional) – the device of the constructed tensor. If None and data is a tensor then the device of data is used. If None and data is not a tensor then the result tensor is constructed on the CPU.

[mksit]

I am sorry that I did not explain clearly. The tensor is created on device=hidden_states.device, but the current creation method has caused non-negligible overheads due to the data copying in my program, which can be seen in the following trace.

This overhead seems unnecessary, so I have suggested this change in this commit.

[fxmarty]

@mksit Are you sure the aten::to op is not there anymore when replacing torch.Tensor by torch.zeros? I generally don't trust pytorch profiler for the timing of aten::to. @NouamaneTazi may have more insights

[mksit]

@fxmarty The aten::to operation disappeared after the replacement in my case. What do you suggest for the profiling of aten::to?

[fxmarty]

After profiling, LGTM. Indeed torch.zeros avoids an aten::to.

[ArthurZucker]

Thanks

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.