I understand your solution - but I guess my problem then is that the function I want to write in terms of PyTorch operations isn’t really re-writable as PyTorch operations (at least from what I have been able to work out).
I have asked a follow-up question here particular to my problem - thanks for your input @KFrank !