Yeah, I’ve tried something similar. But I personally think this can make your model become incredibly bloated.
I think its “easier” to comprehend something like
class Model(nn.Module):

    def __init__(self,args) -> None:
        super(Model).__init__()

        self.args = args

        self.m1 = Module1()

        self.m2 = Module2()

    def forward(x):
        
        match self.args.logic:
            case 'm1':
                x = self.m1(x)
            case 'm2':
                x = self.m2(x)
            case 'm1m2':
                x = self.m1(x)
                x = self.m2(x)
                
        return x

But I’m also thinking the above is not best practise either lol