Update src/model_load.py

src/model_load.py

@@ -12,8 +12,9 @@ def load_model():
     max_memory = {0: "24GB", "cpu": "30GB"}
     # Cargar tokenizer y modelo de Hugging Face
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoModelForCausalLM.from_pretrained(model_name,
-                                                 torch_dtype=torch.bfloat16).to("cuda")
+                                                 torch_dtype=torch.bfloat16).to(device)
     
     # Crear pipeline de generación de texto
     text_generation_pipeline = pipeline(