I think the loss calculation might blow up your memory usage.
Let’s have a look at distMatrix.
You are calling this function with tZ, which has dimensions [25059, 2] and therefore has 50118 elements.
Then you are creating x and y. They have the same shape of [25059, 25059, 2], so 1,255,906,962 elements each.
If we use 4 bytes (float32) for each element, we would have (50118+2*1255906962)*4 / 1e9 ~ 10.05GB memory consumption for this method.
With the model itself you end up with more or less 12GB.
For the training you also need to store some activations, which also take some memory.
How much memory does your GPU has?