In the setup function inside the Module class in gemma.py, should the static_argnums parameter in nn.remat be set to (5,) or (6,)?

[xliu0105]

In the setup function inside the Module class in gemma.py, should the static_argnums parameter in nn.remat be set to (5,) or (6,)? I think it should be the deterministic parameter in the Block class that is set to static — so why is the parameter here set to (5,)?

https://github.com/Physical-Intelligence/openpi/blob/15a9616a00943ada6c20a0f158e3adb39df2ccac/src/openpi/models/gemma.py#L362