Calculating gradient with respect to modified weights TensorFlow

Question

I am trying to implement the Sharpness Aware Minimization (SAM) method in a custom TensorFlow Training Loop. The algorithm follows these steps:

• Calculate gradient w