Fix key and value heads patching for models with different n_heads from n_key_value_heads#981

jlarson4 merged 5 commits intoTransformerLensOrg:devTransformerLensOrg/TransformerLens:devfrom

nikolaystanishev:patching_key_value_heads_bugnikolaystanishev/TransformerLens:patching_key_value_heads_bugCopy head branch name to clipboard