The code looks generally alright, but you’ve switched the model order.
In your initial post you’ve claimed net2 -> net1 -> output, while you are now using net1 -> net2 -> output.
Also, I don’t know which criterion you are using, but often the model output is the first argument and the target the second, so you might want to switch this as well.