Fixing device loading issue

app.py

@@ -17,7 +17,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 loaded_models = {
     "DialoGPT-med-FT": AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 }
-loaded_models["DialoGPT-med-FT"].load_state_dict(torch.load(model_names["DialoGPT-med-FT"]))
+loaded_models["DialoGPT-med-FT"].load_state_dict(torch.load(model_names["DialoGPT-med-FT"], map_location=device))
 loaded_models["DialoGPT-med-FT"].to(device)
 
 loaded_models["DialoGPT-medium"] = AutoModelForCausalLM.from_pretrained(model_names["DialoGPT-medium"]).to(device)
@@ -40,11 +40,11 @@ def respond(
     input_text += f"User: {message}\nAssistant:"
 
     # Tokenize the input text using the shared tokenizer
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
 
     # Generate the response using the selected DialoGPT model
     output_tokens = model.generate(
-        inputs["input_ids"],
+        inputs["input_ids"].to(device),
         max_length=len(inputs["input_ids"][0]) + max_tokens,
         temperature=temperature,
         top_p=top_p,