Aclass(torch.nn.Module) pass the derivations to the inputs

My purpose is to create a class which inherits all methods from torch.nn.Module.

Inside the fit function it would be a specific calculation, I wonder how to deal with derivations? I have a analytical derivation for this function but I don’t know how to extend the class.

Then the parameters of that function are always the same, that means that the training would not effect affect them, that class need to pass the gradient to the inputs during training.

Need your advices guys (examples were great)

Thank you!