Update app.py

app.py

@@ -12,12 +12,11 @@ BASE_MODEL = "TheBloke/stable-vicuna-13B-HF"
 device_map = {0: [0, 1, 2], 1: [3, 4, 5]}
 
 model = LlamaForCausalLM.from_pretrained(
-    BASE_MODEL, 
-    load_in_8bit=True, 
-    load_in_8bit_fp32_cpu_offload=True,
-    torch_dtype=torch.float16, 
-    device_map=device_map,
-    offload_folder="./cache"
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    load_in_8bit=True,
+    device_map="auto",
+    offload_folder="./cache",
 )
 
 tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)