But eventually something needs to have the if statement right?
If I only have this branching code in one module. Eventually the caller also needs to have an if statement for knowing whether it is in training mode
For example
either I write this
ModuleA(jit.ScriptModule):
  @jit.script_method
  def forward(self, inputs):
      self.module_b(inputs)

ModuleB(jit.ScriptModule):
  @jit.script_method 
  def forward(self, inputs):
     if TRAINING:  
         ///
     else: 
        ///

or I write this:
ModuleA(jit.ScriptModule):
  @jit.script_method
  def forward(self, inputs):
      if TRAINING:
         self.module_b.forward_training(inputs)
     else:
         self.module_b.forward_inference(inputs)         


ModuleB(jit.ScriptModule):
   @jit.script_method
   def forward_inference(self, inputs)
   def forward_training(self, inputs):


I don’t want to have to make every single module in my model to have to have forward_training and forward_inference just because one module needs to branch between training and inference