Update app.py

app.py

@@ -1,13 +1,18 @@
 import gradio as gr
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
+import torch
 import spaces
+
 # Define the model and sampling parameters
 model_name = "mistralai/Ministral-8B-Instruct-2410"
 sampling_params = SamplingParams(max_tokens=8192)
 
-# Initialize the LLM model
-llm = LLM(model=model_name, tokenizer_mode="mistral", config_format="mistral", load_format="mistral")
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Initialize the LLM model with the specified device
+llm = LLM(model=model_name, tokenizer_mode="mistral", config_format="mistral", load_format="mistral", device=device)
 
 @spaces.GPU
 # Define the chatbot function
@@ -46,4 +51,4 @@ with gr.Blocks() as demo:
     txt.submit(chatbot, [txt, chatbot], [chatbot, txt])
 
 # Launch the Gradio interface
-demo.launch()
+demo.launch()