I use binary gates self.alpha_gate , which is a tensor composed of 0 and 1, to specify the usage of candidate operations.

Do you save these results as self. in the nn.Module during the forward?? Because you read active index from self.active_index[0] which is not computed here and self.alpha_gate is a different Tensor from what I can see.
In general, you want your nn.Modules to be as stateless as possible. So you might want to avoid saving forward values as self..

I want to save the time of forward computation of all the candidate operations, as there is only one of them contributing to the final output.

Why does it save time given that you do these computations in the backward anyway? You could do them in the forward and just save the result no?