Hi, I’m just learning about PyTorch because I need to migrate my TF 1.x code in order to be able to use adv robustness evaluation libraries. I found the process pretty painless so far, but I’m having problems understanding the difference between the TensorFlow function tf.gradients which computes the gradient vectors of a variable out with respect to a variable in, and the torch.autograd.grad function, which seems to work on scalars only instead. I’ve read other similar entries on this website but I couldn’t find relevant examples to my problem so here I am.
Going into specifics, I need to migrate the following function written in TF 1.x to PyTorch:

def cross_lipschitz(f, x, batch_size=100, z_dim=10):
reg = 0
grad_matrix_list = [tf.gradients(f[:, k], x)[0] for k in range(z_dim)]
for l in range(z_dim):
for m in range(l + 1, z_dim):
grad_diff_matrix = grad_matrix_list[l] - grad_matrix_list[m]
norm_for_batch = tf.norm(grad_diff_matrix, ord=2, axis=1)
reg += tf.reduce_sum(tf.square(norm_for_batch))
return reg

Can you help me out in understanding how to replace tf.gradients with torch.autograd.grad in this particular situation? I also apologize for the wrong code formatting, it’s not clear to me how to do indendation in the forum text editor.