Remember that what works on the forward pass will not necessarily work on the reverse (i.e. taking gradients and backprop) because they are entirely different code. Since you’re training a normal denoising autoencoder and not a VAE, the easiest way to go about things would be to add noise to the inputs before feeding them to the model. This can even be done on the fly by passing train_data through a generator.