Yes, the code looks generally alright.
If your model isn’t training properly, I would recommend to try to overfit a small data snippet (e.g. just 10 samples) and make sure your model is able to perfectly fit this set.
PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier 