Computing gradient of the model with modified weights

Question

I was implementing Sharpness Aware Minimization (SAM) using Tensorflow. The algorithm is simplified as follows

1. Compute gradient using current weight W