testing gpu

app.py

@@ -63,10 +63,11 @@ def quantize_model(model_name, quantization_type, group_size=128, auth_token=Non
     print(f"Quantizing model: {quantization_type}")
     if quantization_type == "int4_weight_only" : 
         quantization_config = TorchAoConfig(quantization_type, group_size=group_size)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
     else : 
         quantization_config = TorchAoConfig(quantization_type)
-    
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
+
     return model
 
 def save_model(model, model_name, quantization_type, group_size=128, username=None, auth_token=None, quantized_model_name=None):
@@ -111,7 +112,7 @@ def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToke
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown(
         """
-        # 🚀 Model Quantization App
+        # 🚀 Model Quantization App :hugging-torch:
         
         Quantize your favorite Hugging Face models and save them to your profile!
         """