state_cached_keys_values append during inference

[Wonder1905]

Hey! I've a small question to discuss, in:
https://github.com/lucidrains/pi-zero-pytorch/blob/main/pi_zero_pytorch/pi_zero.py#L1321C1-L1321C76
We keep appending the state_keys and state_values during inference.
But doesnt the state_keys/values stay "static" after the first "prefill" (still inferencing==False).
I'll try to make myself clearer:
When the diffusion process that ejects the actions is in the diffusion iterations, why would the static language command+image keys and values change?
My question is offcourse only in the current code state where the model only output actions.
Maybe its for the sake of the memory?

Thanks a lot!

[lucidrains]

@Wonder1905 hey! yes, while it isn't broken, i should just directly assign it to the input coming in.

thank you for identifying this issue!

[lucidrains]

@Wonder1905 want to see if the latest commit makes it more clear?

[Wonder1905]

Yes, thanks!