Hello Ergnoor,
great.

Ah, are you implementing orthogonal / unitary RNN by chance?
I did that once, but I didn’t put it in an optimizer.
I would recommend to look at torch.optim.SGD to see how you would implement optimizers in PyTorch.
If you want to do it manually, don’t calculate the grads_ but just take the param.grad from the params after someone called backward.
So you’d drop grad_ = ..., then you replace

for param, grad_ in zip(params, grads_):
    W = param.double()
    G = grad_.double()
    ...
    updates[...] = ...

with
with torch.no_grad(): # we don't actually want autograd in the gradient step
  for param in params:
     W = param.double()
     G = param.grad.double()
...
     param = torch.mm(cayley, W).float() # in stead of updates[param]=...


I’m not sure why one would not be possible or the other. Torch will track your ops all right. You can pass multiple things to differentiate by to grad at the same time or using backward will do the right thing, too. Note that you can only take derivatives of scalar functions.

Best regards
Thomas
P.S.: If you use triple backticks ``` before and after your code, you’ll get all your code formatted. That makes it much nicer to look at.