Hi,
The problem is that the argmax operation that you do on y_pred at the beginning is not differentiable. It returns integer indices.
So no gradient will be able to flow back to y_pred and to your network above