Update app.py

app.py

@@ -65,7 +65,7 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
 
         # Create a visualization
         fig, ax = plt.subplots(figsize=(10, 5))
-        ax.bar(["Original", "Pruned"], [config.num_parameters, pruned_model.num_parameters])
+        ax.bar(["Original", "Pruned"], [config.num_parameters, sum(p.numel() for p in pruned_model.parameters())])
         ax.set_ylabel("Number of Parameters")
         ax.set_title("Model Size Comparison")
         buf = BytesIO()