Loss doesn't convergence

[Xuzhuobin]

Hello, I encountered a situation where the loss doesn't converge while running the code. I wonder if there might be an error in updating the parameter U. My assumption is that when there's a significant error between W and Z, the parameter U rapidly increases and remains high, causing the loss to not converge. Is there any relevant optimization detail that hasn't been updated in the code?