Tensorflow implementation of KdNet