Hi – so I’m defining a torch.autograd.Function.
class MyFunction(torch.autograd.Function):
  @staticmethod
  def forward(ctx, input1, input2):
    # some computation
    ctx.save_for_backward(....)
    return loss
  
  @staticmethod
  def backward(ctx, grad_output):
    # some computation
    return grad1, grad2

From what I understand, anything that happens to input1 inside of forward is not tracked and the gradient applied is whatever is returned in backward. (For example, if loss = input1 * 2, but grad1 = grad_output * 10, the computation graph won’t care about the first statement and treat the gradient as grad_output * 10). Is my understanding correct? Does torch treat the tensor differently because the operation happens inside a torch.autograd.Function rather than a module?
Assuming this understanding is correct, is there a way to combine both behaviors? In particular, suppose my forward function looks something like this:
# inside class MyFunction(torch.autograd.Function)
def forward(ctx, input1, input2):
  a = some_computation_a(input1)
  b = some_computation_b(input1)
  return a + b

All of the steps of some_computation_a aren’t directly differentiable so I want to supply the gradient to those operations myself, but all of the operations in some_computation_b are, so I’d like torch to track those operations automatically. Is this possible to do?
Thanks,
James