Define Function, and transfer into scriptModule, Error

class MultiOp(torch.autograd.Function):
    def forward(ctx, input):
        multiNum = 2
        result = multiNum * input
        return result
    def backward(ctx, grad_output):
        multiNum = ctx.saved_variables
        result = grad_output / multiNum
        return result
multiOp = MultiOp.apply

class Net(torch.jit.ScriptModule):
    def __init__(self):
        super(Net, self).__init__()
    def forward(self, x):
        x = multiOp(x)
        return x
scriptNet = Net(), "")
RuntimeError                              Traceback (most recent call last)
<ipython-input-55-40f7e60ca520> in <module>
----> 1, "")

~/miniconda3/envs/dingyongchao/lib/python3.7/site-packages/torch/jit/ in save(m, f, _extra_files)
    196             (sys.version_info[0] == 2 and isinstance(f, unicode)) or \
    197             (sys.version_info[0] == 3 and isinstance(f, pathlib.Path)):
--> 198, _extra_files=_extra_files)
    199     else:
    200         ret = m.save_to_buffer(_extra_files=_extra_files)

~/miniconda3/envs/dingyongchao/lib/python3.7/site-packages/torch/jit/ in save(self, *args, **kwargs)
   1204         def save(self, *args, **kwargs):
-> 1205             return*args, **kwargs)
   1207         def save_to_buffer(self, *args, **kwargs):

could not export python function call MultiOp. Remove calls to Python functions before export. Did you forget add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__.:
def forward(self, x):
    x = multiOp(x)
        ~~~~~~~ <--- HERE
    return x

So I want to how to define a Function ?

This might be related to @tom’s PR.