Commit 223b121d authored by Frederick Liu's avatar Frederick Liu Committed by A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 401911051
parent 65d1c7ff
...@@ -15,11 +15,11 @@ ...@@ -15,11 +15,11 @@
"""Keras-based TransformerEncoder block layer.""" """Keras-based TransformerEncoder block layer."""
import tensorflow as tf import tensorflow as tf
from official.nlp import keras_nlp from official.nlp import modeling
from official.vision.beta.modeling.layers.nn_layers import StochasticDepth from official.vision.beta.modeling.layers.nn_layers import StochasticDepth
class TransformerEncoderBlock(keras_nlp.layers.TransformerEncoderBlock): class TransformerEncoderBlock(modeling.layers.TransformerEncoderBlock):
"""TransformerEncoderBlock layer with stochastic depth.""" """TransformerEncoderBlock layer with stochastic depth."""
def __init__(self, *args, stochastic_depth_drop_rate=0.0, **kwargs): def __init__(self, *args, stochastic_depth_drop_rate=0.0, **kwargs):
......
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