check for gpu

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from torch import nn
 import gradio as gr
 import os
+import torch
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
 model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
@@ -50,6 +51,9 @@ def annotation(image:ImageDraw, color_seg:np.array):
 
 def call(image): #nparray
 
+    print(f"Is CUDA available: {torch.cuda.is_available()}")
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
     resized = Image.fromarray(image).resize((1024,1024))
     resized_image = np.array(resized)
     print(f"{np.array(resized_image).shape=}") # 1024, 1024, 3