Pytorch multi task

Hello mates,
I have a question, let’s say I have an input image which I would then create many versions of it using autoencoder and multi task learning
Now for testing, I want the user to give let’s say a task number so the output would be only that version of image. Is it possible to do it with pytorch ?
Thanks in advance