--- run_pretraining.py	2020-12-21 17:37:35.000000000 +0100
+++ run_pretraining.py	2021-01-12 01:28:50.000000000 +0100
@@ -143,7 +143,7 @@
 
     self.total_time = 0.0 # total time taken to train (excluding warmup + ckpt saving steps)
     self.step_time = 0.0 # time taken per step
-    self.init_global_step = session.run(tf.train.get_global_step()) # training starts at init_global_step
+    self.init_global_step = session.run(tf.compat.v1.train.get_global_step()) # training starts at init_global_step
     self.skipped = 0
 
   def before_run(self, run_context):
@@ -292,7 +292,7 @@
     total_loss = masked_lm_loss + next_sentence_loss
     total_loss = tf.identity(total_loss, name='total_loss')
 
-    tvars = tf.trainable_variables()
+    tvars = tf.compat.v1.trainable_variables()
 
     initialized_variable_names = {}
     if init_checkpoint and (hvd is None or hvd.rank() == 0):
@@ -380,10 +380,10 @@
   """Get loss and log probs for the masked LM."""
   input_tensor = gather_indexes(input_tensor, positions)
 
-  with tf.variable_scope("cls/predictions"):
+  with tf.compat.v1.variable_scope("cls/predictions"):
     # We apply one more non-linear transformation before the output layer.
     # This matrix is not used after pre-training.
-    with tf.variable_scope("transform"):
+    with tf.compat.v1.variable_scope("transform"):
       input_tensor = tf.layers.dense(
           input_tensor,
           units=bert_config.hidden_size,
@@ -394,7 +394,7 @@
 
     # The output weights are the same as the input embeddings, but there is
     # an output-only bias for each token.
-    output_bias = tf.get_variable(
+    output_bias = tf.compat.v1.get_variable(
         "output_bias",
         shape=[bert_config.vocab_size],
         initializer=tf.zeros_initializer())
@@ -425,12 +425,12 @@
 
   # Simple binary classification. Note that 0 is "next sentence" and 1 is
   # "random sentence". This weight matrix is not used after pre-training.
-  with tf.variable_scope("cls/seq_relationship"):
-    output_weights = tf.get_variable(
+  with tf.compat.v1.variable_scope("cls/seq_relationship"):
+    output_weights = tf.compat.v1.get_variable(
         "output_weights",
         shape=[2, bert_config.hidden_size],
         initializer=modeling.create_initializer(bert_config.initializer_range))
-    output_bias = tf.get_variable(
+    output_bias = tf.compat.v1.get_variable(
         "output_bias", shape=[2], initializer=tf.zeros_initializer())
 
     logits = tf.matmul(tf.cast(input_tensor, tf.float32), output_weights, transpose_b=True)
@@ -537,7 +537,7 @@
   for name in list(example.keys()):
     t = example[name]
     if t.dtype == tf.int64:
-      t = tf.to_int32(t)
+      t = tf.cast(t, 'int32')
     example[name] = t
 
   return example
@@ -591,7 +591,6 @@
       session_config=config,
       save_checkpoints_steps=FLAGS.save_checkpoints_steps if not FLAGS.horovod or hvd.rank() == 0 else None,
       save_summary_steps=FLAGS.save_checkpoints_steps if not FLAGS.horovod or hvd.rank() == 0 else None,
-      keep_checkpoint_max=100,
       # This variable controls how often estimator reports examples/sec.
       # Default value is every 100 steps.
       # When --report_loss is True, we set to very large value to prevent
@@ -602,7 +601,7 @@
   model_fn = model_fn_builder(
       bert_config=bert_config,
       init_checkpoint=FLAGS.init_checkpoint,
-      learning_rate=FLAGS.learning_rate if not FLAGS.horovod else FLAGS.learning_rate*hvd.size(),
+      learning_rate=FLAGS.learning_rate,
       num_train_steps=FLAGS.num_train_steps,
       num_warmup_steps=FLAGS.num_warmup_steps,
       use_one_hot_embeddings=False,
