RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Hello @mkang315 , I was trying to use this code base for my project. When I use the train script I get the following error:

File "/teamspace/studios/this_studio/CST-YOLO/utils/loss.py", line 759, in build_targets
    from_which_layer = from_which_layer[fg_mask_inboxes]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[mkang315]

The error message you’re seeing, “RuntimeError: indices should be either on CPU or on the same device as the indexed tensor (CPU),” typically occurs in PyTorch or a similar deep learning framework. This error arises when you’re trying to perform an operation that involves tensors on different devices (e.g., CPU and GPU) and the indices used for indexing are not on the same device as the tensor you’re trying to index.

To resolve this issue, you should ensure that the indices you’re using for indexing are on the same device as the tensor you’re trying to index. Here’s how you can do it:

Check Device Compatibility: First, make sure that the tensor you’re trying to index and the indices you’re using are compatible in terms of device. If you’re using a GPU tensor, the indices should also be on the same GPU.

Move Indices to the Correct Device: If your indices are on the wrong device, you can move them to the correct device using the .to(device) method. For example, if you have a GPU tensor gpu_tensor and you want to move the indices to the GPU.

Our code is based on YOLOv7. Please refer to link1 and link2 in the YOLOv7 repo.