If you can do it with torchtensors then you can define your own custom autograd function.
It might be a little tedious, but we’ve all written our own custom functions and modules - that’s a lot of the fun  .
It’s definitely worth doing, once you’ve done one, it’s a lot less scary - good luck