Update app.py

app.py

@@ -13,7 +13,12 @@ model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it",
 tok = AutoTokenizer.from_pretrained("google/gemma-2b-it",token=token)
 # using CUDA for an optimal experience
 # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-device = torch.device('cuda')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
+else:
+    device = torch.device('cpu')
+    print("Using CPU")
 model = model.to(device)