diff --git a/keras_nlp/models/xlnet/relative_attention.py b/keras_nlp/models/xlnet/relative_attention.py index d11dc4bd5a..a11ae3fd9d 100644 --- a/keras_nlp/models/xlnet/relative_attention.py +++ b/keras_nlp/models/xlnet/relative_attention.py @@ -15,7 +15,6 @@ import math import string -from keras_nlp.api_export import keras_nlp_export from keras_nlp.backend import keras from keras_nlp.backend import ops @@ -76,7 +75,6 @@ def _rel_shift(x, klen=-1): return x -@keras_nlp_export("keras_nlp.layers.TwoStreamRelativeAttention") class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention): """Two-stream relative self-attention for XLNet. diff --git a/keras_nlp/models/xlnet/xlnet_encoder.py b/keras_nlp/models/xlnet/xlnet_encoder.py index 13f5a953ee..bb8e56e4cc 100644 --- a/keras_nlp/models/xlnet/xlnet_encoder.py +++ b/keras_nlp/models/xlnet/xlnet_encoder.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from keras_nlp.api_export import keras_nlp_export from keras_nlp.backend import keras from keras_nlp.backend import ops from keras_nlp.models.xlnet.relative_attention import TwoStreamRelativeAttention @@ -22,7 +21,6 @@ def xlnet_kernel_initializer(stddev=0.02): return keras.initializers.TruncatedNormal(stddev=stddev) -@keras_nlp_export("keras_nlp.layers.XLNetEncoder") class XLNetEncoder(keras.layers.Layer): """ XLNet Encoder.