@rasbt @ptrblck Thanks for the explanation. The code I am modifying is from the MNIST example on this website 44.
I also tried the example 210 from the official website of PyTorch, it is fast and converge to good performance. Even if I remove the dropout part, and modify the neural network as the tensorflow example, I get better results (around 95%).
I will focus on the official example from PyTorch. Thanks for your explanations.