|
bboxes = pred[0].boxes.xyxy.cpu().numpy() |
I've noticed that there numpy used in places. Could we have a torch check to see if device is cuda, and to use torch instead of numpy when a gpu is available? This will greatly speed things up instead of using the cpu in these instances.
PyTorch supports both CPU and GPU acceleration, so an if statement probably wouldn't even be necessary.
asdff/asdff/yolo.py
Line 71 in 8ac7245
I've noticed that there numpy used in places. Could we have a torch check to see if device is cuda, and to use torch instead of numpy when a gpu is available? This will greatly speed things up instead of using the cpu in these instances.
PyTorch supports both CPU and GPU acceleration, so an if statement probably wouldn't even be necessary.