Couldn't reproduce the MNISTSuperPixel results

[dokato]

I tried to reproduce your accuracies from the paper, but running the vanilla model straight from the repo( without any tweaks), doesn't look like its learning:

Epoch: 001, Loss: 6.97393, Train Acc: 0.10218, Test Acc: 0.10100
Epoch: 002, Loss: 2.30191, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 003, Loss: 2.30144, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 004, Loss: 2.30132, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 005, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 006, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 007, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 008, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 009, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 010, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 011, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 012, Loss: 2.30125, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 013, Loss: 2.30821, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 014, Loss: 2.30259, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 015, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 016, Loss: 2.30351, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 017, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 018, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 019, Loss: 2.30125, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 020, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 021, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 022, Loss: 2.30125, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 023, Loss: 2.30119, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 024, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 025, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 026, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 027, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 028, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 029, Loss: 2.30232, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 030, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 031, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 032, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 033, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 034, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 035, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 036, Loss: 2.30123, Train Acc: 0.11237, Test Acc: 0.11350
Epoch: 037, Loss: 2.30124, Train Acc: 0.11237, Test Acc: 0.11350

I tried with and without data augmentations and varied learning rate as well, but no success in fixing it.