diff --git a/tensorflow/python/keras/engine/distributed_training_utils.py b/tensorflow/python/keras/engine/distributed_training_utils.py
index cd8e0342cdd..c587287ff96 100644
--- a/tensorflow/python/keras/engine/distributed_training_utils.py
+++ b/tensorflow/python/keras/engine/distributed_training_utils.py
@@ -328,7 +328,7 @@ def init_restore_or_wait_for_variables():
   """Initialize or restore variables or wait for variables to be initialized."""
   session = K._get_session()  # pylint: disable=protected-access
   worker_context = dc_context.get_current_worker_context()
-  if not worker_context or worker_context.should_init:
+  if not worker_context or worker_context.experimental_should_init:
     # TODO(yuefengz): if checkpoints exit, restore from checkpoint.
     K._initialize_variables(session)  # pylint: disable=protected-access
   else: