b3377b098a78e8a2508527f9d71a89b15fef3929,official/nlp/modeling/layers/on_device_embedding.py,OnDeviceEmbedding,call,#OnDeviceEmbedding#Any#,80

Before Change

      embeddings = tf.matmul(one_hot_data, self.embeddings)
    else:
      embeddings = tf.gather(self.embeddings, flat_inputs)
    embeddings = tf.reshape(embeddings, input_shape)

    return embeddings

After Change

      embeddings = tf.matmul(one_hot_data, self.embeddings)
    else:
      embeddings = tf.gather(self.embeddings, flat_inputs)
    embeddings = tf.reshape(
        embeddings,
        // Work around b/142213824: prefer concat to shape over a Python list.
        tf.concat([tf.shape(inputs), [self._embedding_width]], axis=0))
    embeddings.set_shape(inputs.shape.as_list() + [self._embedding_width])
    return embeddings

Instances