We discussed this internally, and one other issue that came up was the size of your linear layer. You’re doing a sum of 50,000 products, which can accumulate a lot of floating point error (which can differ between different platforms). You should get better results with a more traditional model that uses some convolutional and pooling layers before the linear layer.