Unable to load the YOLOv5 on to the GPU

[yasakova-anastasia]

My actions before raising this issue

• Read/searched the docs
• Searched past issues

I have been trying to integrate the model yolov5m trained on Objects365 - https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5m_Objects365.pt onto the GPU on CVAT

Expected Behaviour

The custom YOLOv5 model to be accessed on the CVAT and be used for annotation using GPU.

Current Behaviour

The building fails and unable to load onto the GPU on CVAT. When deploying on CPU, it works. But not sure of the changes to be made to deploy using GPU

Steps to Reproduce (for bugs)

• function-gpu.yaml -
  Changed the base image which was configured for CUDA version 10.2
  

• function-gpu.yaml
  Added the pip install requirements and the torch and torchvision and download the model.
  

The requirements are from here - https://github.com/ultralytics/yolov5/blob/master/requirements.txt

• function-gpu.yaml
  The resources for GPU and changed the maxWorkers for triggers to 1 as read in one of the previous issues.
  

• main.py
  Access the downloaded weights file and load to the available cuda device (in my case, 1 GPU) and continue with inference.
  

Context

I am unable to resolve this issue no matter the changes I make in the main.py file or the function-gpu.yaml.

Your Environment

• Docker version docker version : 20.10.17

• Operating System and version - Linux

• Code example or link to GitHub repo to access the model : https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5m_Objects365.pt

• Other diagnostic information / logs:

Thank you for your kind help. Looking forward to any help/guidelines to solve the issue.

Issue by @gd1925