Thank you a lot for the reply!
Which is the best practice currently suggested for cases similar to the one of the original poster?
Let’s say I want to implement a module with a forward of the type:
def forward(self, x):
   if torch.allclose(x,self.known_input):
      return self.know_output
   else:
     return self.mlp(x)

Is there a simple way to keep that conditional computation inside the module, just yielding zero gradients when doing the backward pass through it?