Hi Ruotian,
Sorry I just posted the core part of the code. The flag is randomly generated to be either True or False.
            flag = random.choice([True, False])

And yes, it goes to both the “reconstruction” part and the “prediction” part of the code. Is this type of code supported by PyTorch? I noticed that there is an example code using the similar idea https://github.com/jcjohnson/pytorch-examples#pytorch-control-flow–weight-sharing 22.
Thank you for your help.
Eric