Update app.py

app.py

@@ -12,9 +12,6 @@ from transformers import (
     LlamaTokenizer,
 )
 
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 MAX_MAX_NEW_TOKENS = 1024
 DEFAULT_MAX_NEW_TOKENS = 50
 MAX_INPUT_TOKEN_LENGTH = 512
@@ -29,7 +26,9 @@ For additional detail on the model, including a link to the arXiv paper, refer t
 
 model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct", 
-    trust_remote_code=True, 
+    device_map="cuda", 
+    torch_dtype="auto", 
+    trust_remote_code=True,
 )
 tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")