[BugFix] Enhance test_pos_encoding to support execution on multi-devices #13187
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
"The test case in test_pos_encoding.py frequently fails when run in a multi-device environment.
FAILED test_pos_encoding.py::test_rotary_embedding[cuda:1-0-dtype0-32-120-17-8192-5-False] - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument index in method wrapper_CUDA__inde...
The issue arises because the rope created in this case does not specify a device, so it is always created on device 0, even if a default device is set. This patch fixes the problem, allowing the test case to pass when running on multiple devices.