Update app.py

app.py

@@ -14,14 +14,16 @@ model_id = "meta-llama/Llama-2-7b-chat-hf"
 
 quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
 device_map = {
-    "transformer.word_embeddings": 0,
-    "transformer.word_embeddings_layernorm": 0,
-    "lm_head": "cpu",
-    "transformer.h": 0,
-    "transformer.ln_f": 0,
+"transformer.word_embeddings": "cpu",
+"transformer.word_embeddings_layernorm": "cpu",
+"lm_head": "cpu",
+"transformer.h": "cpu",
+"transformer.ln_f": "cpu",
+"model.embed_tokens": "cpu",
+"model.layers":"cpu",
+"model.norm":"cpu"
 }
 
-
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map,quantization_config=quantization_config)
 tokenizer = AutoTokenizer.from_pretrained(model_id)