tiny fix

app.py

@@ -15,13 +15,13 @@ class GradioApp:
 
     def __init__(self) -> None:
 
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-        custom = CustomUnet().to(device).eval()
-        custom.load_state_dict(torch.load('models/custom_unet.pt', map_location=device))
+        custom = CustomUnet().to(self.device).eval()
+        custom.load_state_dict(torch.load('models/custom_unet.pt', map_location=self.device))
         
-        pretrained = get_pretrained_unet().to(device).eval()
-        pretrained.load_state_dict(torch.load('models/pretrained_unet.pt', map_location=device))
+        pretrained = get_pretrained_unet().to(self.device).eval()
+        pretrained.load_state_dict(torch.load('models/pretrained_unet.pt', map_location=self.device))
 
         self.models = {
             'Custom': custom,