I want to use conditional computation of the following form. It uses the output of NN1 in stage 1, in order to decide whether to use NN2 or NN3 in the next stage.
def forward(self, x):
      y = self.NN1(x)
      if y>0:
         return self.NN2(x)
     else:
         return self.NN3(x)

I do not know how to implement it. The existing methods compute both NN2(x) and NN3(x), and then multiply either NN2(x) by 0, or NN3(x) by 0. However, such an implementation does not give savings in terms of computations saved.