Hi,

Is my understanding correct?

Yes you are absolutely correct!

Does torch treat the tensor differently because the operation happens inside a torch.autograd.Function rather than a module?

No, the Tensors are treated the same way, but everything runs with autograd disable diring the forward.

Is this possible to do?

Of course!
I think the simplest solution is to (and is also a general guideline) keep your custom Function as small as possible: It will make the gradients easier to derive for you.
In particular, you can have a custom function that handles fw/bw for some_computation_a only and then combine it is a regular python function:
def my_fn(input1, input2):
    return SomeComputationA.apply(input1) + some_computation_b(input2)

Does that work for you?