Thank you so much, that helped a lot! @soulitzer  I had another doubt, my grads for A will be the same size of A. So if my A is of size 10 then grads will of size 1x10 I have another set of grads that I calculated wrt to the model’s parameters and hence is a tuple of length 512 containing 2-d tensors. (Eg: ((tensor of size 10x10), (tensor of size 10), …) → dummy sizes) Now my task is to multiply the grads i.e A.grads x model.grads and use the resultant to update A. Any idea how to do so?