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)
    elif cond == 1:
      return self.B(input)
    else:
      return self.C(input)

Now whether such model will work well for your task I don’t know. I’m afraid you will have to try  and see experimentally if it does.