Implementation of the gradient approximation for the SAM objective (eq3)

I have read your paper and code carefully.
You use the p as 2, thus SAM firstly maximizes the objective loss function by adding rho * normalized gradients to models' parameters (https://github.com/google-research/sam/blob/dae9904c4cf3a57a304f7b04cecffe371679c702/sam_jax/training_utils/flax_training.py#L537)
And SAM calculates gradients in this noised model again (https://github.com/google-research/sam/blob/dae9904c4cf3a57a304f7b04cecffe371679c702/sam_jax/training_utils/flax_training.py#L549).
In your paper, the gradient is calculated with respect to the original model's parameters, unlike the code which calculates gradient with respect to the noised model's parameters.

Are these two the same?