Access python defined global variable in c++ autograd.function


I need to code a Jit compatible custom operation where the backward-pass depends on the value of a global variable defined in my python code.
Is there a way to do that?

In practice, that would correspond to the following piece of code in python

FOO = True

class MyFunnyFun(Funtion):
    def forward(ctx, x):
        # return something

    def backward(ctx, grad_output):
        if FOO:
            # return something
            # return something else



jit does not actually support custom Function no?

If you don’t need the jit, this code will work just fine with regular pytorch :slight_smile: