Update app.py

app.py

@@ -1,21 +1,21 @@
-import gradio as gr
-import torch
+import os
+import ssl
 import cv2
+import torch
+import certifi
 import numpy as np
-import torchvision.transforms as transforms
-from PIL import Image, ImageEnhance
-from torchvision import models
+import gradio as gr
 import torch.nn as nn
-import matplotlib.pyplot as plt
+from torchvision import models
 import torch.nn.functional as F
-import ssl
-import certifi
-import os
+import matplotlib.pyplot as plt
+from PIL import Image, ImageEnhance
+import torchvision.transforms as transforms
 
 ssl._create_default_https_context = lambda: ssl.create_default_context(cafile=certifi.where())
 
 # Set device
-device = "cpu"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Number of classes
 num_classes = 6