Update app.py

app.py

@@ -6,7 +6,20 @@ from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
 from transformers import Trainer
 
 BASE_MODEL = "TheBloke/stable-vicuna-13B-HF"
-model = LlamaForCausalLM.from_pretrained(BASE_MODEL, load_in_8bit=True, torch_dtype=torch.float16, device_map="auto", offload_folder="./cache")
+
+# Create a custom device map
+# This will vary based on the architecture of your model and the memory capacity of your GPU and CPU
+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"
+)
+
 tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token_id = 0
 tokenizer.padding_side = "left"