In your code, you seem to be only manipulating only parts of the keys and values
all_store_kv['key'].append(self.scale * key[:, :, 512:, :][:, :, dim_neg2_indices, :])
all_store_kv['value'].append(value[:, :, 512:, :][:, :, dim_neg2_indices, :])Is there a reason you do this? Why don't you use the all of the key or value?
In your code, you seem to be only manipulating only parts of the keys and values
Is there a reason you do this? Why don't you use the all of the key or value?