Different forward( ) in train and test modes

Dear All,

I want to implement a two-branch architecture with forward( ) operate differently in train and test mode:

  • For training, forward(x1, x2) which allows model takes x1, x2 for branch_A and branch_B at the same time.

  • For testing, I want to evaluate each branch separately, such as forward(x1) and forward(x2).

How to implement this?

Thank you.

You can use the internal self.training flag in a condition and then use the desired code path.
Other layers, such as nn.Dropout and nn.BatchNormXd, are doing the same.
The self.training attribute will be set via model.train() and model.eval().

1 Like