ds

app.py

@@ -42,15 +42,16 @@ def validate_parameters(max_tokens, temperature, top_p):
         return False, "Error: 'Top-p' must be between 0.1 and 1.0."
     return True, ""
 
+# Determine the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load the model and tokenizer
 model_name = "gpt2"  # Use GPT-2 model
 
 try:
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
-    model = transformers.AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-    )
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
+    model.to(device)
     model.eval()
 except Exception as e:
     logging.error(f"Failed to load model {model_name}: {e}")
@@ -76,7 +77,7 @@ def respond(message, history, persona_choice, custom_persona, max_tokens, temper
     logging.info(f"Received message: {safe_message}")
 
     try:
-        input_ids = tokenizer.encode(conversation, return_tensors="pt").to(model.device)
+        input_ids = tokenizer.encode(conversation, return_tensors="pt").to(device)
 
         output_ids = model.generate(
             input_ids,