fix: wrong grid device w.r.t. output in yolox head

[pierrot-lc]

Hi,

When using your amazing repo, I had to fix a small issue where the devices were wrong because some tensor is not using the same device as the model.

The issue can be found here:

def get_output_and_grid(self, output, k, stride, dtype):
    grid = self.grids[k]

    batch_size = output.shape[0]
    n_ch = 5 + self.num_classes
    hsize, wsize = output.shape[-2:]
    if grid.shape[2:4] != output.shape[2:4]:
        yv, xv = meshgrid([torch.arange(hsize), torch.arange(wsize)])
        grid = torch.stack((xv, yv), 2).view(1, 1, hsize, wsize, 2).type(dtype)
        self.grids[k] = grid

    output = output.view(batch_size, 1, n_ch, hsize, wsize)
    output = output.permute(0, 1, 3, 4, 2).reshape(
        batch_size, hsize * wsize, -1
    )
    grid = grid.view(1, -1, 2)
    output[..., :2] = (output[..., :2] + grid) * stride  # <- Here grid and output can be on different devices.
    output[..., 2:4] = torch.exp(output[..., 2:4]) * stride
    return output, grid

A simple fix could be:

def get_output_and_grid(self, output, k, stride, dtype):
    ...
    grid = grid.view(1, -1, 2).to(output.device)  # <- To make sure the devices are the same.
    output[..., :2] = (output[..., :2] + grid) * stride
    output[..., 2:4] = torch.exp(output[..., 2:4]) * stride
    return output, grid

I'm not sure why I seem to be the only one having this issue, I hope I'm not using the model the wrong way. Note that my use case is a little bit complicated so I had to implement the training loops myself (that may explain why I faced this issue).

Thank you for your time and have a great day.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

Pierre Pereira seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}