cuda test

app/app.py

@@ -1,6 +1,7 @@
 from flask import Flask
 from doclayout_yolo import YOLOv10
 from huggingface_hub import hf_hub_download
+import torch
 
 import os
 import json
@@ -35,11 +36,13 @@ def doc_layout():
         if file and allowed_file(file.filename):
             filename = "/tmp/" + secure_filename(file.filename)
             file.save(filename)
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            print(device)
             det_res = model.predict(
                 filename,   # Image to predict
                 imgsz=1024,        # Prediction image size
                 conf=0.2,          # Confidence threshold
-                device="cuda:0"    # Device to use (e.g., 'cuda:0' or 'cpu')
+                device=device
             )
             names = det_res[0].names
             blocknames = [names[int(n)] for n in det_res[0].boxes.cls]