MNIST Autoencoder training

1

Heyy!
Here’s a situation:
I am using an Autoencoder on MNIST dataset. Firstly, I trained the AE with MSE loss function as:

  1. Encoder- input layer(784,300), hidden layer(300,64)
  2. Decoder- hidden layer(64,300), output layer(300,784)
    I used leaky relu as activation function.
    Now I am taking the encoder portion and attaching an output layer of (64,10) dimensions and training it using log_softmax activation in the output layer and target as the digit the data represents, for example, an array containing [1], [4], etc. Here, I am using nll loss function.
    Now, when I apply this network on test data, error is very high.

Any solution to propose?