Update app.py

app.py

@@ -7,15 +7,11 @@ from transformers import Trainer
 
 BASE_MODEL = "TheBloke/vicuna-7B-1.1-HF"
 
-# Create a custom device map
-# This will vary based on the architecture of model and the memory capacity of GPU and CPU
-device_map = {0: [0, 1, 2], 1: [3, 4, 5]}
-
 model = LlamaForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
     load_in_8bit=True,
-    device_map = {0: [0, 1, 2], 1: [3, 4, 5]},
+    device_map = "auto",
     offload_folder="./cache",
 )