Indices not moved to the correct device in src/trackformer/models/detr_tracking.py` line 96:

[TNodeCode]

I found a small bug in the code. When I train trackformer on my own custom dataset I get the following error:

Traceback (most recent call last):
  File "src/train.py", line 356, in <module>
    train(args)
  File "src/train.py", line 283, in train
    visualizers['train'], args)
  File "/scratch/dldevel/flasche/trackformer/src/trackformer/engine.py", line 126, in train_one_epoch
    outputs, targets, *_ = model(samples, targets)
  File "/scratch/dldevel/flasche/miniconda3/envs/trackformer/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/scratch/dldevel/flasche/trackformer/src/trackformer/models/detr_tracking.py", line 262, in forward
    self.add_track_queries_to_targets(targets, prev_indices, prev_out)
  File "/scratch/dldevel/flasche/trackformer/src/trackformer/models/detr_tracking.py", line 100, in add_track_queries_to_targets
    for ind in not_prev_out_ind
  File "/scratch/dldevel/flasche/trackformer/src/trackformer/models/detr_tracking.py", line 101, in <listcomp>
    if ind not in prev_out_ind]
  File "/scratch/dldevel/flasche/miniconda3/envs/trackformer/lib/python3.7/site-packages/torch/tensor.py", line 659, in __contains__
    return (element == self).any().item()  # type: ignore[union-attr]
  File "/scratch/dldevel/flasche/miniconda3/envs/trackformer/lib/python3.7/site-packages/torch/tensor.py", line 27, in wrapped
    return f(*args, **kwargs)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

It is caused by this statement in src/trackformer/models/detr_tracking.py line 96:

not_prev_out_ind = torch.arange(prev_out['pred_boxes'].shape[1])

I could fix the error by moving not_prev_out_ind to the used device:

not_prev_out_ind = torch.arange(prev_out['pred_boxes'].shape[1]).to(device)