diff --git a/tf_keras/export/export_lib.py b/tf_keras/export/export_lib.py
index 88a967cfc..4e44fc753 100644
--- a/tf_keras/export/export_lib.py
+++ b/tf_keras/export/export_lib.py
@@ -244,7 +244,7 @@ def serving_fn(x):
             decorated_fn = tf.function(fn, input_signature=input_signature)
             self._endpoint_signatures[name] = input_signature
         else:
-            if isinstance(fn, tf.types.experimental.GenericFunction):
+            if isinstance(fn, tf.types.experimental.PolymorphicFunction):
                 if not fn._list_all_concrete_functions():
                     raise ValueError(
                         f"The provided tf.function '{fn}' "