Thanks for the quick response! On implementing a custom function, I did have a look at creating custom autograd functions which have forward and backward methods. I was wondering if it were at all possible to implement a backward of backward method in an autograd function?
What I mean by backward of backward would be an analytical expression for the 2nd order gradient of the loss function with respect to the output of the final layer. Because the final layer of my network is a determinant function, explicitly writing out those derivatives and calling them directly would be far more efficient than calling autograd and letting it solve it itself.