internal change

PiperOrigin-RevId: 334829363

official/nlp/keras_nlp/layers/masked_lm.py

@@ -81,8 +81,7 @@ class MaskedLM(tf.keras.layers.Layer):
     lm_data = self.layer_norm(lm_data)
     lm_data = tf.matmul(lm_data, self.embedding_table, transpose_b=True)
     logits = tf.nn.bias_add(lm_data, self.bias)
-
-    masked_positions_length = masked_positions.shape[1] or tf.shape(
+    masked_positions_length = masked_positions.shape.as_list()[1] or tf.shape(
         masked_positions)[1]
     logits = tf.reshape(logits,
                         [-1, masked_positions_length, self._vocab_size])