Solved by ptrblck in post #2 
              
              
                You could split your autoencoder into the encoder and decoder part and then train all modules together. 
Here is a dummy example: 
encoder = MyEncoder()
decoder = MyDecoder()

z = encoder(data)
out = decoder(z)

loss_z = criterion_z(z, target_z)
loss_out = criterion_out(out, target_out)
loss = loss_…