Replace dynamic Squeeze with reshape to 1d

TPUs don't support dynamic Squeeze as it won't be able to identify the dimensions to remove if the input shape is dynamic. For example, [<=16, 1] could return a vector or scalar result depending on the size of the dynamic dim.

PiperOrigin-RevId: 415083667

official/nlp/bert/run_classifier.py

@@ -76,7 +76,7 @@ def get_loss_fn(num_classes):
 
   def classification_loss_fn(labels, logits):
     """Classification loss."""
-    labels = tf.squeeze(labels)
+    labels = tf.reshape(labels, [-1])
     log_probs = tf.nn.log_softmax(logits, axis=-1)
     one_hot_labels = tf.one_hot(
         tf.cast(labels, dtype=tf.int32), depth=num_classes, dtype=tf.float32)