From 18f4b9274f1e20182047e2cc312b4ba3ed7b61bc Mon Sep 17 00:00:00 2001 From: Gunnlaugur Thor Briem Date: Wed, 4 Mar 2020 16:57:28 +0000 Subject: [PATCH] fix: work with Tensorflow < 2.1.0 tf.keras.utils.register_keras_serializable was added in TF 2.1.0, so don't rely on it being there; just decorate the class with it if it exists. --- src/transformers/modeling_tf_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py index e10afb37a8..3c94c6bf60 100644 --- a/src/transformers/modeling_tf_utils.py +++ b/src/transformers/modeling_tf_utils.py @@ -72,7 +72,9 @@ def keras_serializable(cls): cls.get_config = get_config cls._keras_serializable = True - return tf.keras.utils.register_keras_serializable()(cls) + if hasattr(tf.keras.utils, "register_keras_serializable"): + cls = tf.keras.utils.register_keras_serializable()(cls) + return cls class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin):