hyunwookim:

how can we check that the gradients (or partial derivatives in this case precisely) are calculated for functions outside of the linear layer?


I am not sure I understand this question.
What is the function in the code above for which you want to compute gradients, and with respect to which variables do you want to compute the gradients of this function?