Adding class label to Autoencoder

Hey PyTorch community,

I am using an Autoenoder for reconstruction purposes. My current problem is that I want to take the class label into account. So, during the training Autoenoder knows that which type of data (class) is using by inputting a class label (or probability of multiple classes). During the testing, by having that label class (or probability of multiple classes) the Autoenoder can reconstruct but based on the class pre-knowledge. I do not know how to implement the network or loss function in this case. Would you please help me?