Update app.py

app.py

@@ -2,7 +2,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import gradio as gr
 import os
-
+device = "cuda" if torch.cuda.is_available() else "cpu"
+prin(device)
 # Asegúrate de que tu token de Hugging Face está cargado como una variable de entorno
 hf_token = os.environ.get("token")
 
@@ -24,7 +25,7 @@ model_base = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
 model_base.resize_token_embeddings(len(tokenizer))
 
 # Cargar el modelo con el adaptador
-model_with_adapter = PeftModel.from_pretrained(model_base, "Juliofc/chaterapia_model")
+model_with_adapter = PeftModel.from_pretrained(model_base, "Juliofc/chaterapia_model").to(device)
 
 
 # Suponiendo que `tokenizer` y `model_with_adapter` ya están inicializados