Your training code looks fine to me. A few things to try out:
- Check how good your training data is and what’s the datasize like. A larger amount of high quality (and unbiased dataset) would result in a better performance.
- Once, you’re happy with the initial model, try different hyperparams (number of conv layers, hidden size, learning rate etc).
- Run for more number of epochs. Calculate the average loss for each epoch (not just each iteration).