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

nikolaystanishev wants to merge 2 commits intoTransformerLensOrg:mainTransformerLensOrg/TransformerLens:mainfrom

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