[WIP] Fixes for Nearest Neighbor

Fixing up some things left over from the merge of the two labs.