Got it. Refer Course 4 Week 4 project 1 - #2 by dbiber
Need to specify the axis when we reduce_sum for distances.