You haven’t specified n_h here. Can you check to see if its value is not too large? If n_h is comparable to n_x, model may just learn to memorize entire input data and not generalize. If the model is overfitting and you don’t have enough data for validation set, try using smaller n_h. Alternatively you could do K-fold cross validation to avoid creating separate validation set. Using train-validation loss plot would give you the exact idea about when to stop training to avoid overfitting. Also, you have defined dropout but don’t seem to be using it.