Tracing custom autograd.Function

jit
1

I have a custom autograd.Function:

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

    @staticmethod
    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\__init__.py", line 903, in trace
    name = _qualified_name(func)
  File "...\Python37\lib\site-packages\torch\_jit_internal.py", 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.

2

Reproduced. Could you please open a new issue at https://github.com/pytorch/pytorch/issues?

3

Opened it: https://github.com/pytorch/pytorch/issues/32822

1 Like