diff --git a/lm/modeling.py b/lm/modeling.py
index 7e9e5330..ccf69504 100644
--- a/lm/modeling.py
+++ b/lm/modeling.py
@@ -18,7 +18,9 @@
 import math
 
 import six
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 from lm import optimization_adafactor
 from lm.utils import get_assignment_map_from_checkpoint, get_shape_list, get_attention_mask, gelu, layer_norm, dropout, \
@@ -84,7 +86,7 @@ def from_dict(cls, json_object):
     @classmethod
     def from_json_file(cls, json_file):
         """Constructs a `NewsConfig` from a json file of parameters."""
-        with tf.gfile.GFile(json_file, "r") as reader:
+        with tf.io.gfile.GFile(json_file, "r") as reader:
             text = reader.read()
         return cls.from_dict(json.loads(text))
 
diff --git a/lm/optimization_adafactor.py b/lm/optimization_adafactor.py
index b8d03ed1..954e112f 100644
--- a/lm/optimization_adafactor.py
+++ b/lm/optimization_adafactor.py
@@ -85,7 +85,7 @@ def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
     return train_op, train_metrics
 
 
-class AdaFactorOptimizer(tf.train.Optimizer):
+class AdaFactorOptimizer(tf.compat.v1.train.Optimizer):
     """here's the optimizer we'll use"""
 
     def __init__(self,
diff --git a/lm/utils.py b/lm/utils.py
index aa75c71b..bc971825 100644
--- a/lm/utils.py
+++ b/lm/utils.py
@@ -17,7 +17,9 @@
 import re
 
 import six
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 import numpy as np
 from tensorflow.python.lib.io import file_io