Port pytorch code to tf2.0: equivalent of x_batch.requires_grad = True in tf2.0?

Question

I am trying to port pytorch code to tf2 code for this repo. The overall logic of the code is as follows:

• sample data from Gaussian distribution with zero mean a