thanks a lot for the reply…but my problem is that the next layers (which are linear1 and linear2)  might not be linear at all…similarly, the condition I apply, might also not be linear…general case if as follows
def forward(self, x):
    if self.NN1(x)>0:
        return self.NN2(x)
   else:
        return self.NN3(x)

Basically, for data samples x that have NN1(x)>0, I want to output NN2(NN1(x)), and for data points x that have NN1(x)<0, NN3(NN1(x)). I am wondering whether your solution will cover this case. Thanks for the reply!