Update app.py

app.py

@@ -15,7 +15,7 @@ model = LlamaForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
     load_in_8bit=True,
-    device_map="auto",
+    device_map = {0: [0, 1, 2], 1: [3, 4, 5]},
     offload_folder="./cache",
 )