Hi,
First of all you do not need to explicitly register the parameters. This is done for you automatically when you save an nn.Parameter as self.foo. This means that you can remove your register_hook function completely and it won’t change anything.
The problem here is that the operations that you do with self.mybranch is not differentiable. Indeed the gradient will be 0 almost everywhere and infinite at the point where it’s equal to 1.
What you can do is to do a “smoother” version of this branching:
def forward(self, x):
   out = self.conv1(x)
   out = self.conv1_relu(out)
   out = self.conv2(out)
   out = self.conv2_relu(out)
   # Always compute the branch
   branch = self.conv3(out)
   branch = self.conv3_relu(branch)
   # Smooth choise between out and branch
   #  make choise variable between 0 and 1
   branch_choice = torch.sigmoid(self.mybranch)
   out = branch_choice * branch + (1 - branch_choice) * out
   out = self.linear1(out)
return out