Hi,
You can find an example here: Gradient zero for one parameter in custom model resulting in no update