We currently don’t support autograd.Function in Python. Right now the workaround is to define the function in C++ and bind it to TorchScript as a custom op 38. This was done in pytorch/vision to support the Mask R-CNN model, you can see the specific implementation here 47.