You can use the cpp api equivalent: torch::autograd::Function. And you will need to bind the ::apply() method properly to python.
EDIT: You can find here 1 an example of the cpp Function. You can then bind its apply method just like any function.