Update app.py

app.py

@@ -7,6 +7,7 @@ import requests
 import torch
 import datetime
 import subprocess
+import spaces  # Ensure this import is correct and the module is available
 
 CUSTOM_CSS = """
 #output_box textarea {
@@ -24,7 +25,8 @@ if not os.path.exists(model_path):
         f.write(response.content)
 
 # Load the document segmentation model
-docseg_model = YOLO(model_path)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+docseg_model = YOLO(model_path).to(device)
 
 def process_image(image):
     # Convert image to the format YOLO model expects
@@ -43,7 +45,7 @@ def process_image(image):
     return annotated_img, detected_areas_labels
 
 zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+print(zero.device) # <-- 'cuda:0' if GPU is available, otherwise 'cpu'
 
 @spaces.GPU
 def run_gpu() -> str: