Fix quantization test for XLNet (#1699)

keras-team · Jul 19, 2024 · f5c1fc1 · f5c1fc1
1 parent b0c21b3
commit f5c1fc1
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/keras_nlp/src/models/xlnet/xlnet_backbone_test.py b/keras_nlp/src/models/xlnet/xlnet_backbone_test.py
@@ -40,7 +40,6 @@ def test_backbone_basics(self):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
             expected_output_shape=(2, 5, 2),
-            run_quantization_check=False,  # TODO(hongyu): set to `True`
         )
 
     @pytest.mark.large

diff --git a/keras_nlp/src/models/xlnet/xlnet_encoder.py b/keras_nlp/src/models/xlnet/xlnet_encoder.py
@@ -100,6 +100,7 @@ def build(self, input_shape):
             dtype=self.dtype_policy,
             name="rel_attn",
         )
+        self.relative_attention.build(input_shape)
 
         self.layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,