change to dynamic device

app.py

@@ -5,7 +5,10 @@ import gradio as gr
 import spaces
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-device = "cuda" # for GPU usage or "cpu" for CPU usage
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained("PatronusAI/Llama-3-Patronus-Lynx-8B-Instruct")
 model = AutoModelForCausalLM.from_pretrained("PatronusAI/Llama-3-Patronus-Lynx-8B-Instruct", torch_dtype=torch.float16, device_map="auto")