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

Hi,

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):
    @staticmethod
    def forward(ctx, x):
        # return something

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

Thanks!

Hi,

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: