Hi, I tried to train a yolov4 pytorch model with a reference to this repo:
  
      
      github.com
  
  
    

WongKinYiu/PyTorch_YOLOv4 2


  PyTorch implementation of YOLOv4. Contribute to WongKinYiu/PyTorch_YOLOv4 development by creating an account on GitHub.

  
  
    
    
  
  


Then I run into a error like this
Traceback (most recent call last):
  File "train.py", line 438, in <module>
    train(hyp, opt, device, tb_writer)
  File "train.py", line 301, in train
    results, maps, times = test.test(opt.data,
  File "/home/kamal/Documents/PyTorch_YOLOv4/test.py", line 206, in test
    plot_images(img, output_to_target(output, width, height), paths, str(f), names)  # predictions
  File "/home/kamal/Documents/PyTorch_YOLOv4/utils/general.py", line 907, in output_to_target
    return np.array(targets)
  File "/home/kamal/.virtualenvs/pytorch-yolov4/lib/python3.8/site-packages/torch/tensor.py", line 621, in __array__
    return self.numpy()
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

My pytorch version is 1.8.0 and cuda is 11.2
Can somebody look into this?