Update app.py

app.py

@@ -45,20 +45,20 @@ def load_model(model_name, progress=gr.Progress()):
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 torch_dtype=torch.float16,
-                device_map="cuda:0",
+                device_map="auto",
                 load_in_8bit=True
             )
         elif "llama" in model_name.lower() or "mistral" in model_name.lower():
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 torch_dtype=torch.float16,
-                device_map="cuda:0"
+                device_map="cpu"
             )
         else:
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 torch_dtype=torch.float16,
-                device_map="cuda:0"
+                device_map="cpu"
             )
         
         if tokenizer.pad_token is None: