Tracing custom autograd.Function

I have a custom autograd.Function:

class raw(autograd.Function):
    def forward(ctx, inp):
        ctx.a = (inp * inp + 1).reciprocal()
        ctx.b = ctx.a.sqrt()
        return inp * ctx.b

    def backward(ctx, grad_output):
        return grad_output * ctx.a * ctx.b

and I want to trace this function, but when I do:

jit.trace(raw.apply, example_inputs=tc.randn(1))

, I get the error from this line:

    jit.trace(raw.apply, example_inputs=tc.randn(1))
  File "...\Python37\lib\site-packages\torch\jit\", line 903, in trace
    name = _qualified_name(func)
  File "...\Python37\lib\site-packages\torch\", line 696, in _qualified_name
    "__module__ can't be None.".format(name))
RuntimeError: Could not get qualified name for class 'apply': __module__ can't be None.

This code used to work for pytorch 1.1.0, but after I updated it to 1.4.0 recently, I got this error. Iā€™m using Python3.7.3 on windows10.

Reproduced. Could you please open a new issue at

Opened it:

1 Like