Hi Mahdi!



 mahdi_morafah:

I want to freeze W1, W2 and only update a, and b through the training.


Let me assume that W1, W2, a, and b are some sort of pytorch
tensors (and that your “.” represents some sort of multiplication).
Simply set
W1.requires_grad = False
W2.requires_grad = False
a.requires_grad = True
b.requires_grad = True

When you backpropagate, pytorch will only calculate gradients
for a and b, treating W1 and W2 as fixed (not trained) parameters.
Best.
K. Frank