Wrong device when tracing models

We are trying to use tracing to export detectron2 models, and are running into a problem in the ImageList.from_tensor method, previously reported https://github.com/facebookresearch/detectron2/issues/4580 and https://github.com/triton-inference-server/server/issues/2024. The problem is caused by the max_sizes tensor, created from the shapes_to_tensor function. Although this function accepts a device as an argument, this argument is not used in the method, so the tensor is returned on CPU. However, if size_divisibility > 1, the max_size is modified using some constants at line 101. In eager mode, these are just constants and there's no problem, but when tracing, these are tensors that live on GPU, resulting in a device mis-match.

I have a two-line fix to this problem which I would be happy to contribute - the tensors just need to be moved to the same device. However, I am a US government employee, and there can be issues with licenses related to work produced by US government employees, since our work product cannot be covered by copyright. Per guidance, I am inquiring prior to submitting a pull request.

Instructions To Reproduce the Bug:

1. Full runnable code or full changes you made:

Attempted to export a trained model using tracing.

1. What exact command you run:

Attempt to trace the model and then run a network through the model.

1. Full logs or other relevant observations:

The traced model fails on ImageList.from_tensor, due to a wrong device error.

1. please simplify the steps as much as possible so they do not require additional resources to run, such as a private dataset.

Expected behavior:

N/A

Environment:

N/A