Update app.py

app.py

@@ -10,6 +10,11 @@ from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
 
 tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
 
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
 model = LlamaForCausalLM.from_pretrained(
     "decapoda-research/llama-7b-hf",
     load_in_8bit=True,