Solved by KaneFu in post #3 
              
              
                Updated
I found that in Keras, random_normal initializer initalize the weights with Normal(0, 0.05) not Normal(0, 1). So I changed the pytorch last layer weights initialization to 
nn.init.normal_(param.data, std=std)  results becomes better. I was surprised that in such a simple task, weights initi…