--- optimization.py	2020-11-04 01:04:32.000000000 +0100
+++ optimization.py	2020-12-21 20:41:32.000000000 +0100
@@ -106,20 +106,20 @@
                                                                                     decr_ratio=0.5)
     optimizer = tf.contrib.mixed_precision.LossScaleOptimizer(optimizer, loss_scale_manager)
 
-  tvars = tf.trainable_variables()
+  tvars = tf.compat.v1.trainable_variables()
   grads_and_vars = optimizer.compute_gradients(loss * 1.0 / num_accumulation_steps, tvars)
 
   if num_accumulation_steps > 1:
-      local_step = tf.get_variable(name="local_step", shape=[], dtype=tf.int32, trainable=False,
+      local_step = tf.compat.v1.get_variable(name="local_step", shape=[], dtype=tf.int32, trainable=False,
                                    initializer=tf.zeros_initializer)
-      batch_finite = tf.get_variable(name="batch_finite", shape=[], dtype=tf.bool, trainable=False,
+      batch_finite = tf.compat.v1.get_variable(name="batch_finite", shape=[], dtype=tf.bool, trainable=False,
                                      initializer=tf.ones_initializer)
-      accum_vars = [tf.get_variable(
+      accum_vars = [tf.compat.v1.get_variable(
           name=tvar.name.split(":")[0] + "/accum",
           shape=tvar.shape.as_list(),
           dtype=tf.float32,
           trainable=False,
-          initializer=tf.zeros_initializer()) for tvar in tf.trainable_variables()]
+          initializer=tf.zeros_initializer()) for tvar in tf.compat.v1.trainable_variables()]
 
       reset_step = tf.cast(tf.math.equal(local_step % num_accumulation_steps, 0), dtype=tf.bool)
       local_step = tf.cond(reset_step, lambda:local_step.assign(tf.ones_like(local_step)), lambda:local_step.assign_add(1))
@@ -127,7 +127,7 @@
       grads_and_vars_and_accums = [(gv[0],gv[1],accum_vars[i]) for i, gv in enumerate(grads_and_vars) if gv[0] is not None]
       grads, tvars, accum_vars = list(zip(*grads_and_vars_and_accums))
 
-      all_are_finite = tf.reduce_all([tf.reduce_all(tf.is_finite(g)) for g in grads]) if manual_fp16 or use_fp16 else tf.constant(True, dtype=tf.bool)
+      all_are_finite = tf.reduce_all([tf.reduce_all(tf.math.is_finite(g)) for g in grads]) if manual_fp16 or use_fp16 else tf.constant(True, dtype=tf.bool)
       batch_finite = tf.cond(reset_step,
         lambda: batch_finite.assign(tf.math.logical_and(tf.constant(True, dtype=tf.bool), all_are_finite)),
         lambda:batch_finite.assign(tf.math.logical_and(batch_finite, all_are_finite)))
@@ -139,7 +139,7 @@
             grads, clip_norm=1.0,
             use_norm=tf.cond(
                 all_are_finite,
-                lambda: tf.global_norm(grads),
+                lambda: tf.linalg.global_norm(grads),
                 lambda: tf.constant(1.0)))
 
       accum_vars = tf.cond(reset_step,
@@ -166,7 +166,7 @@
       grads_and_vars = [(g, v) for g, v in grads_and_vars if g is not None]
       grads, tvars = list(zip(*grads_and_vars))
       all_are_finite = tf.reduce_all(
-          [tf.reduce_all(tf.is_finite(g)) for g in grads]) if use_fp16 or manual_fp16 else tf.constant(True, dtype=tf.bool)
+          [tf.reduce_all(tf.math.is_finite(g)) for g in grads]) if use_fp16 or manual_fp16 else tf.constant(True, dtype=tf.bool)
 
       # This is how the model was pre-trained.
       # ensure global norm is a finite number
@@ -175,7 +175,7 @@
           grads, clip_norm=1.0,
           use_norm=tf.cond(
               all_are_finite,
-              lambda: tf.global_norm(grads),
+              lambda: tf.linalg.global_norm(grads),
               lambda: tf.constant(1.0)))
 
       train_op = optimizer.apply_gradients(
@@ -220,7 +220,7 @@
       has_shadow = manual_fp16 and param.dtype.base_dtype != tf.float32
       if has_shadow:
         # create shadow fp32 weights for fp16 variable
-        param_fp32 = tf.get_variable(
+        param_fp32 = tf.compat.v1.get_variable(
             name=param_name + "/shadow",
             dtype=tf.float32,
             trainable=False,
@@ -228,13 +228,13 @@
       else:
         param_fp32 = param
 
-      m = tf.get_variable(
+      m = tf.compat.v1.get_variable(
           name=param_name + "/adam_m",
           shape=param.shape.as_list(),
           dtype=tf.float32,
           trainable=False,
           initializer=tf.zeros_initializer())
-      v = tf.get_variable(
+      v = tf.compat.v1.get_variable(
           name=param_name + "/adam_v",
           shape=param.shape.as_list(),
           dtype=tf.float32,
@@ -325,7 +325,7 @@
       has_shadow = manual_fp16 and param.dtype.base_dtype != tf.float32
       if has_shadow:
         # create shadow fp32 weights for fp16 variable
-        param_fp32 = tf.get_variable(
+        param_fp32 = tf.compat.v1.get_variable(
             name=param_name + "/shadow",
             dtype=tf.float32,
             trainable=False,
@@ -333,13 +333,13 @@
       else:
         param_fp32 = param
 
-      m = tf.get_variable(
+      m = tf.compat.v1.get_variable(
           name=param_name + "/adam_m",
           shape=param.shape.as_list(),
           dtype=tf.float32,
           trainable=False,
           initializer=tf.zeros_initializer())
-      v = tf.get_variable(
+      v = tf.compat.v1.get_variable(
           name=param_name + "/adam_v",
           shape=param.shape.as_list(),
           dtype=tf.float32,
