Problem
I am trying to develop an architecture that haves some computation that is only triggered by some external signal (for example, when the loss between two epochs does not change much).
However, I am not sure how I would implement this because of two reasons

If I write a conditional in the model definition that looks like following. I am not sure this conditional will be executed

class NeuralNet(nn.Module):
  # ...
  def forward(self, x):
    if <conditional>:
      # do something
    x = ...
    return x


I could not use hook function to get this done since my operation will affect both forward and backward pass, but there is a known bug of register_backward_hook.

Any input is appreciated!