Update src/model_load.py

src/model_load.py

@@ -14,10 +14,7 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name,
                                                  torch_dtype=torch.float16,
-                                                 device_map={
-                                                     "transformer.h.0": "cpu",  # Por ejemplo, mover ciertas capas al CPU
-                                                     "transformer.h.1": "cuda",  # Mantener otras capas en la GPU
-                                                 },
+                                                 device_map="auto"
                                                  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
                                                  llm_int8_enable_fp32_cpu_offload=True)