Thanks for the update. It seems your model’s parameters get valid gradients.
I would recommend to play around with some hyperparameters, such as lowering the learning rate.
If that doesn’t help, try to scale down the problem and try to overfit a small data sample (e.g. only 10 samples). Once your model is able to do so, you could try to scale it up again carefully.