Thanks for the answer.
After changes you suggested, gradients are computed correctly for every element in the matrix. This cause another problem. After update of weights the matrix is no longer a rotation matrix (orthogonal and with determinant equal to 1). I want to remain its structure so that after each iteration it would be still a rotation matrix, just like in constructor. So I’d like to update only parameters a, b, c.
In addition, next question came up. Where do I restrict a, b, c to stay in the [-1,1] segment to force these parameters to act like cosine functions? I thought of something like:
rotation_matrix.a = torch.max(torch.tensor(-1.), torch.min(a, torch.tensor(1.)))

but have no idea what place in the code is good for it.