Skip to content

GitLab

Commit a6e44930 authored by Hongkun Yu's avatar Hongkun Yu Committed by A. Unique TensorFlower
Browse files

Nit: allow use runtime shape for hub module without explicit set_shape. 

PiperOrigin-RevId: 288423138
parent 1e1333fe
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/nlp/bert_models.py
View file @ a6e44930
...@@ -256,7 +256,8 @@ class BertSquadLogitsLayer(tf.keras.layers.Layer): ...@@ -256,7 +256,8 @@ class BertSquadLogitsLayer(tf.keras.layers.Layer):
"""Implements call() for the layer.""" """Implements call() for the layer."""
sequence_output = inputs sequence_output = inputs
input_shape = sequence_output.shape.as_list() input_shape = tf_utils.get_shape_list(
sequence_output, name='sequence_output_tensor')
sequence_length = input_shape[1] sequence_length = input_shape[1]
num_hidden_units = input_shape[2] num_hidden_units = input_shape[2]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment