Solved by albanD in post #4 
              
              
                You can have a model with conditionals with no problem in pytorch: 
class MyMod(nn.Module):
  def __init__(self, params):
    self.A = MyModA()
    self.B = MyModB()
    self.C = MyModC()

  def forward(self, input):
    cond = compute_condition(input)
    if cond == 0:
      return self.A(input)
  …