You could try to overfit a small dataset first (e.g. just 10 examples) by playing around with some hyperparameters and make sure your model is able to do so. Once this is done, you could try to scale up the use case again. In case the overfitting approach is not working, your training script might have some bugs (which are not shown in the posted code snippet) or the general approach might not work with the proposed architecture etc.