Updated app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
@@ -11,7 +12,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     # "kroonen/phi-2-GGUF",
     "microsoft/phi-2",
-    # device_map="cpu",
+    device_map="auto",
     trust_remote_code=True,
     torch_dtype=torch.float32
 )