add cuda information retrieval + a variable to force using cuda for ONNX

Dockerfile

@@ -13,6 +13,6 @@ WORKDIR /app
 
 COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
-
+RUN ONNX_MODE=cuda
 COPY --chown=user . /app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

tasks/image.py

@@ -28,6 +28,14 @@ print(12)
 router = APIRouter()
 print(13)
 
+import torch
+
+# Get CUDA version (the one PyTorch was compiled with)
+print("CUDA version:", torch.version.cuda)
+
+# Get cuDNN version
+print("cuDNN version:", torch.backends.cudnn.version())
+
 #MODEL_TYPE = "YOLOv11n"
 DESCRIPTION = f"YOLOv8n model with batch 128 inference on CPU"
 print(14)