I wouldn’t worry about this too much performance-wise, but note that there is a way to omit runtime checks with JIT:

class Model(Module):
  flag : jit.Final[bool]
  def __init__(self, out_features):
     ...
     self.flag = out_features > 1
     self.tail = Flatten() if self.flag else None
  def forward(self,x):
     ...
     if self.flag:
       x=self.tail(x)
m = jit.script(Model(1))

this should totally remove ‘if’ code block. Even simpler, (tail is not None) check would also be statically resolved.