Fix test breakages after changing assertion message for checkpoint status check.
Fix test breakages after changing assertion message for checkpoint status check.

PiperOrigin-RevId: 408738920

tensorflow/hub,tensorflow/hub