Per my last message, I would remove all uses of params and instead follow the docs on how to take gradients of models. If you’re still running into issues then a MWE and full stacktrace will be required, per Please read: make it easier to help you.
1 Like