Update app.py

app.py

@@ -541,10 +541,12 @@ import traceback
 
 def get_device():
     if torch.cuda.is_available():
-        print('Using CUDA GPU')
-        return torch.device('cuda')
+        device = torch.device('cuda')
+        print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+        print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+        return device
     else:
-        print('Using CPU')
+        print("Using CPU")
         return torch.device('cpu')
 
 device = get_device()