When training with the Mahalanobis distance instead of the Euclidean distance it seems that the network is not training i.e. the loss is not decreasing but remain constant.
I have tried changing the learning rate (both higher and lower LR) but it does not seem to change.
The function may not be well implemented and the gradients might get lost somewhere in the mahalanobis_distance function.