autodiff for user script functions aka torch.jit.script for autograd.Function #22329
Description
I got an error when use jit.script on some new layers(implemented in c++):
RuntimeError:
attribute lookup is not defined on python value of type 'FunctionMeta':
@torch.jit.script_method
def forward(self, x):
result = TopPoolFunction.apply(x)
~~~~~~~~~~~~~~~~~~~~~ <--- HERE
return result
here is my code:
class TopPoolFunction(Function):
@staticmethod
def forward(ctx, input):
output = top_pool.forward(input)[0]
ctx.save_for_backward(input)
return output
@staticmethod
def backward(ctx, grad_output):
input = ctx.saved_variables[0]
output = top_pool.backward(input, grad_output)[0]
return output
class TopPool(torch.jit.ScriptModule):
@torch.jit.script_method
def forward(self, x):
result = TopPoolFunction.apply(x)
return result
What is the proper way to use script on new layers, HELPPPP!