Inference stage issue

[pluto32]

I understand that during the inference stage, a single past_key_values is being used for both the augment model and the anchor model. I'm curious whether this is correct or if I've misunderstood something, because the hidden state values used in the augment model and those used in the anchor model would be different. I'm leaving this question to clarify my understanding.

[lyconghk]

I also have a similar problem. After finishing training CALM, in the inference stage, can I use the default generate function to generate the output for a given input question? Or should I write a new function generate to override the generate function?