About the Operation of the Newly Added Visualize

[Zengyf-CVer]

@glenn-jocher

I just ran the newly added feature map visualization command:

python detect.py --visualize --weights ./weight/yolov5s.pt --source ./train2017/000000000081.jpg

Then this error was reported in the shell:

TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

The following is a specific screenshot:

My configuration:

• python 3.8
• CUDA 11.1
• PyTorch 1.9.0

[glenn-jocher]

@Zengyf-CVer ah, thanks for the bug report.

[glenn-jocher]

@Zengyf-CVer good news 😃! Your original issue may now be fixed ✅ in PR #3925. To receive this update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload with model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!