So the new way is just very new (we added this to TorchVision in September). They probably branched off the maskrcnn-benchmark or TorchVision earlier than that and didn’t get the update.
Previously, you had the problem that “JIT means CustomOp means implemented in C++” and that meant no autograd with jumping through lots of hoops 1. In August or so, we gained C++ autograd::Function, so now you can easily have CustomOps that record a custom backward.
This still isn’t the “ultimate JIT-friendly thing”, which would be more like graph-to-graph differentiation 1, that sadly I didn’t get to finish up yet (I’ll see if adding a voting mechanism for prioritizing my PyTorch work 6 helps…).
Best regards
Thomas