Fix case with 1 prediction.
Fix case with 1 prediction.

Changed tf.squeeze to tf.reshape to avoid getting the wrong shape when having a single element.

tryolabs/luminoth,tryolabs/luminoth,tryolabs/luminoth