Update app.py

app.py

@@ -36,17 +36,18 @@ try:
     )
     logger.info("Tokenizer loaded successfully")
 
-    # Load model
+    # Load model - Modificada la carga del modelo
     logger.info("Loading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        device_map="auto",
         trust_remote_code=True
     )
+    if device == "cuda":
+        model = model.to(device)
     logger.info("Model loaded successfully")
 
-    # Create pipeline
+    # Create pipeline - Modificada la creación del pipeline
     logger.info("Creating generation pipeline...")
     model_gen = pipeline(
         "text-generation",
@@ -57,7 +58,7 @@ try:
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.1,
-        device_map="auto"
+        device=0 if device == "cuda" else -1
     )
     logger.info("Pipeline created successfully")