test zero gpu

app.py

@@ -61,14 +61,22 @@ model = AutoModel.from_pretrained("{model_name}")"""
     return model_card
 
 @spaces.GPU
+def load_model_gpu(model_name, quantization_config, auth_token) : 
+    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
+
+def load_model_cpu(model_name, quantization_config, auth_token) : 
+    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
+
 def quantize_model(model_name, quantization_type, group_size=128, auth_token=None, username=None, device="cuda"):
     print(f"Quantizing model: {quantization_type}")
     if quantization_type == "int4_weight_only" : 
         quantization_config = TorchAoConfig(quantization_type, group_size=group_size)
     else : 
         quantization_config = TorchAoConfig(quantization_type)
-    
-    model = AutoModel.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)
+    if device == "cuda" : 
+        model = load_model_gpu(model_name, quantization_config, auth_token)
+    else : 
+        model = load_model_cpu(model_name, quantization_config, auth_token)
 
     return model
 
@@ -110,11 +118,11 @@ def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToke
         return exists_message
     if quantization_type == "int4_weight_only" and device == "cpu" : 
         return "int4_weight_only not supported on cpu"
-    try : 
-        quantized_model = quantize_model(model_name, quantization_type, group_size, oauth_token, profile.username, device)
-        return save_model(quantized_model, model_name, quantization_type, group_size, profile.username, oauth_token, quantized_model_name)
-    except Exception as e : 
-        return e
+    # try : 
+    quantized_model = quantize_model(model_name, quantization_type, group_size, oauth_token, profile.username, device)
+    return save_model(quantized_model, model_name, quantization_type, group_size, profile.username, oauth_token, quantized_model_name)
+    # except Exception as e : 
+    #     return e
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as app: