Morphological Distance Transform on Pytorch

I want to define function “scipy.ndimage.morphology.distance_transform_edt” in tensor computation.
It can be replaced by using “torch.from_numpy()” and above function, but it’s too slow for traning state.
Is there any implementation or replacement for this computation?

The visualization of function process
image

shortcut for Scipy document “distance_transform_edt”
https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.ndimage.morphology.distance_transform_edt.html

it has been long time, were you able to find the solution ? if so can you share it please.
Thank you.