Update app.py

app.py

@@ -24,6 +24,13 @@ if not hf_token:
 dtype = torch.bfloat16
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    n_gpu = torch.cuda.device_count()
+    torch.cuda.get_device_name(0)
+else:
+    device = torch.device("cpu")
+
 count0 = torch.zeros(1).to(device)
 count1 = torch.zeros(1).to(device)
 count2 = torch.zeros(1).to(device)