Fix to add output module for autoencoder