Update app.py

app.py

@@ -90,6 +90,12 @@ def interact(user_input, history, interaction_count, model_name):
         if tokenizer is None or model is None:
             raise ValueError("Tokenizer or model is not initialized.")
         
+        # Determine the device to use (either CUDA if available, or CPU)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Ensure the model is on the correct device
+        model.to(device)
+        
         if interaction_count >= MAX_INTERACTIONS:
             user_input += ". Thank you for your questions. Our session is now over. Goodbye!"
         
@@ -102,11 +108,8 @@ def interact(user_input, history, interaction_count, model_name):
         
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         
-        # Generate response using selected model
-        input_ids = tokenizer(prompt, return_tensors='pt').input_ids
-        if model_name not in quantized_models:
-            input_ids.to("cuda")
-            
+        # Move input tensor to the same device as the model
+        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)
         chat_history_ids = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id, temperature=0.1)
         response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)