Fix/data: handle models with specific `num_key_value_heads`

[nickfraser]

For some models, num_attention_heads != num_key_value_heads. This fixes the dimension of the past_key_values in this case.