Error when pruning Llama3.2 3B

[parkjiwon1]

I ran the example code for pruning the Llama3.2 3B model provided by examples/LLMs/prune_llm.py but this error occurs.

"RuntimeError: shape '[24, -1]' is invalid for input of size 1024"

As far as I understand, there are 1024 out_features channels of k_proj in the self-attention layer of the Llama 3.2-3B model, which seems to be an error caused by not exactly dividing this number by the group number 24.
But the example code said that it ran successfully, so I wonder how you guys did it.

Here are my environment settings.

torch 2.5.1+cu121
torch-pruning 1.6.0
transformers 4.55.0

Thanks