Update app.py

app.py

@@ -32,8 +32,13 @@ torch_dtype = torch.bfloat16  # Set the appropriate torch data type
 
 # Display device and CPU thread information
 device = get_device()
-print("Running on device:", device)
-print("CPU threads:", torch.get_num_threads())
+print(f"Number of GPUs available: {torch.cuda.device_count()}")
+print(f"Running on device: {device}")
+print(f"CPU threads: {torch.get_num_threads()}")
+
+# Check if CUDA is available and set the device accordingly
+if not torch.cuda.is_available():
+    raise RuntimeError("CUDA is not available. Ensure that a GPU is available and properly configured.")
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)