Custom Model pruning

[saidineshpola]

How to fix this error for custom model pruning? The pruning was done but some layers weren't pruned properly. I got this error at the time of inference

key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) RuntimeError: shape '[1, 7, 4, 128]' is invalid for input of size 2688

[xiaoxuanNLP]

I have the same error