Creating a Custom KS-test Loss Function

Hi all!

I’m developing a neural network as a part of my PhD project and am having trouble choosing the correct loss function. My degree is in nuclear engineering and I am trying to unfold neutron spectra. My inputs for network are a detector response (think individual data points) and outputs are a neutron spectrum (similar to a discretized pdf). From what I’ve read, the Kolmogorov-Smirnov test is the best way to compare if two pdfs come from the same distribution. I’m not sure how to incorporate this test as the loss function using the built in TensorFlow functions. Could someone help point me in the direction to learn?

I hope you have a great day!