diff --git a/generate.py b/generate.py
index ecfd2bf6d..e1bf54894 100644
--- a/generate.py
+++ b/generate.py
@@ -7,11 +7,16 @@
 import os
 
 import librosa
+import soundfile
 import numpy as np
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from wavenet import WaveNetModel, mu_law_decode, mu_law_encode, audio_reader
 
+tf.disable_v2_behavior()
+
+
 SAMPLES = 16000
 TEMPERATURE = 1.0
 LOGDIR = './logdir'
@@ -112,7 +117,7 @@ def _ensure_positive_float(f):
 
 def write_wav(waveform, sample_rate, filename):
     y = np.array(waveform)
-    librosa.output.write_wav(filename, y, sample_rate)
+    soundfile.write(filename, y, sample_rate)
     print('Updated wav file at {}'.format(filename))
 
 
diff --git a/test/test_causal_conv.py b/test/test_causal_conv.py
index 4801ed460..099cb5c1f 100644
--- a/test/test_causal_conv.py
+++ b/test/test_causal_conv.py
@@ -1,7 +1,8 @@
 """Unit tests for the causal_conv op."""
 
 import numpy as np
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from wavenet import time_to_batch, batch_to_time, causal_conv
 
diff --git a/train.py b/train.py
index 02ec80074..932fcf18f 100644
--- a/train.py
+++ b/train.py
@@ -14,11 +14,15 @@
 import sys
 import time
 
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.client import timeline
 
 from wavenet import WaveNetModel, AudioReader, optimizer_factory
 
+tf.disable_v2_behavior()
+
+
 BATCH_SIZE = 1
 DATA_DIRECTORY = './VCTK-Corpus'
 LOGDIR_ROOT = './logdir'
diff --git a/wavenet/audio_reader.py b/wavenet/audio_reader.py
index a65e9c894..fa708d84d 100644
--- a/wavenet/audio_reader.py
+++ b/wavenet/audio_reader.py
@@ -6,7 +6,8 @@
 
 import librosa
 import numpy as np
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 FILE_PATTERN = r'p([0-9]+)_([0-9]+)\.wav'
 
@@ -65,7 +66,7 @@ def trim_silence(audio, threshold, frame_length=2048):
     '''Removes silence at the beginning and end of a sample.'''
     if audio.size < frame_length:
         frame_length = audio.size
-    energy = librosa.feature.rmse(audio, frame_length=frame_length)
+    energy = librosa.feature.rms(y=audio, frame_length=frame_length)
     frames = np.nonzero(energy > threshold)
     indices = librosa.core.frames_to_samples(frames)[1]
 
diff --git a/wavenet/model.py b/wavenet/model.py
index d2e79a9af..6bb5f8219 100644
--- a/wavenet/model.py
+++ b/wavenet/model.py
@@ -1,5 +1,7 @@
 import numpy as np
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
+import tensorflow.keras as keras
 
 from .ops import causal_conv, mu_law_encode
 
@@ -7,7 +9,7 @@
 def create_variable(name, shape):
     '''Create a convolution filter variable with the specified name and shape,
     and initialize it using Xavier initialition.'''
-    initializer = tf.contrib.layers.xavier_initializer_conv2d()
+    initializer = keras.initializers.glorot_normal()
     variable = tf.Variable(initializer(shape=shape), name=name)
     return variable
 
diff --git a/wavenet/ops.py b/wavenet/ops.py
index 304791ba0..726c5158f 100644
--- a/wavenet/ops.py
+++ b/wavenet/ops.py
@@ -1,6 +1,7 @@
 from __future__ import division
 
-import tensorflow as tf
+#import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def create_adam_optimizer(learning_rate, momentum):