Thanks for the reply.
Yes, you are right, but when i apply script on TopPoolFunction, i got an error:
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 how i script TopPool:
class TopPool(torch.jit.ScriptModule):
     # here is what i added, is there any problem?
    @torch.jit.script_method        
    def forward(self, x):
        result = TopPoolFunction.apply(x)
        return result

For TopPoolFunction, i actually have no idea how to script…