If the calculations are “trainable”, that is there’s some learning involved, perhaps you could use a simple multi-layer perception mlp which takes in the output of your model, as input.
The output of the mlp could then feed into the loss criterion.
Maybe posting some code of the calculations you want do might be helpful to understand what you want to do?