Maybe an existing thread and this blog post help you.
I don’t see how you have changed the forward call ? At the moment your network’s forward return the output logits. If you wish you can return x just before classifier as well.