@driazati thank you for your timely answer.
It seems I can only write forward function in current torchscript custom ops, right? What if users would like to customize backward functions like what we did in torch.autograd.Function?