device of unit_locations should follow tensor_input

[aryopg]

Description

In modeling_utils.py, there is a part of the code that causes error when using GPU. It seems that it's because the unit_locations variable was forced to be in cpu.

Testing Done

No test done yet, just to flag the issue and to raise the attention of the main contributors

Checklist:

• [ ] My PR title strictly follows the format: [Your Priority] Your Title
• [ ] I have attached the testing log above
• [ ] I provide enough comments to my code
• [ ] I have changed documentations
• [ ] I have added tests for my changes

[frankaging]

@PinetreePantry can you take a look at this one? thanks!