Oh, I think that would really be an interesting way to approach the problem, I’ll give it a try for sure and update you.
I think, as you say, the models are indeed equivalent. I did some further inspection and realised that I wasn’t using .zero_grad() on the classification layer separately when training. As in, I always called, model.net.zero_grad() followed by loss.backward() and optimizer.step() and this led to the difference in performance. It is quite interesting though to see that not calling model.classification_layer.zero_grad() caused such a large difference. I’ll close this discussion for now and mark it resolved. Thank you for all the help.