Hi, I want two step training:
step 1: train only model_A
step 2: train model_B in connection to model_A that is in eval mode
-Do you think I still need to wrap the forward pass for step2? how to wrap it only for step 2…?
-And how should I understand the situation that the error occurs only when device=cuda ?