Update app.py

app.py

@@ -52,11 +52,10 @@ def load_model(selected_model_name):
     # Load the model with quantization settings
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
+        quantization_config=bnb_config,
         trust_remote_code=True,
     )
 
-    # Ensure every part of the model is assigned to the correct device
-    model.to(device)  # This should correctly set devices for all components
 
     # Additional configurations and training enhancements
     model.config.use_cache = False