Internal change

PiperOrigin-RevId: 430740048

Internal change
PiperOrigin-RevId: 430740048
a0b548e2 · Jiayu Ye · A. Unique TensorFlower · b72f4975 · a0b548e2
Commit a0b548e2 authored Feb 24, 2022 by Jiayu Ye Committed by A. Unique TensorFlower Feb 24, 2022
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

official/nlp/modeling/networks/funnel_transformer.py official/nlp/modeling/networks/funnel_transformer.py +5 -0

No files found.
--- a/official/nlp/modeling/networks/funnel_transformer.py
+++ b/official/nlp/modeling/networks/funnel_transformer.py
@@ -402,6 +402,11 @@ class FunnelTransformerEncoder(tf.keras.layers.Layer):
            _transformer_cls2str.get(transformer_cls, str(transformer_cls))
    }

+    self.inputs = dict(
+        input_word_ids=tf.keras.Input(shape=(None,), dtype=tf.int32),
+        input_mask=tf.keras.Input(shape=(None,), dtype=tf.int32),
+        input_type_ids=tf.keras.Input(shape=(None,), dtype=tf.int32))
+
  def call(self, inputs):
    # inputs are [word_ids, mask, type_ids]
    if isinstance(inputs, (list, tuple)):