diff --git a/docqa/nn/layers.py b/docqa/nn/layers.py
index 7e78ab4..ea00222 100644
--- a/docqa/nn/layers.py
+++ b/docqa/nn/layers.py
@@ -3,8 +3,8 @@
 import numpy as np
 import tensorflow as tf
 from docqa.configurable import Configurable
-from tensorflow.contrib.keras import activations
-from tensorflow.contrib.keras import initializers
+from tensorflow.keras import activations
+from tensorflow.keras import initializers
 from tensorflow.python.layers.core import fully_connected
 
 from docqa.model import Prediction
@@ -802,4 +802,4 @@ def apply(self, is_train, tensor1, tensor2, mask) -> tf.Tensor:
         if self.include_tiled:
             out.append(tf.tile(tf.expand_dims(tensor2, 1), [1, tf.shape(tensor1)[1], 1]))
 
-        return tf.concat(out, axis=2)
\ No newline at end of file
+        return tf.concat(out, axis=2)
diff --git a/docqa/nn/span_prediction.py b/docqa/nn/span_prediction.py
index 838ee07..43c5945 100644
--- a/docqa/nn/span_prediction.py
+++ b/docqa/nn/span_prediction.py
@@ -7,7 +7,7 @@
 from docqa.nn.span_prediction_ops import best_span_from_bounds, to_unpacked_coordinates, \
     to_packed_coordinates, packed_span_f1_mask
 from tensorflow import Tensor
-from tensorflow.contrib.layers import fully_connected
+from tensorflow.keras.layers import Dense  # For TensorFlow 2.x
 
 from docqa.model import Prediction
 from docqa.nn.ops import VERY_NEGATIVE_NUMBER, exp_mask, segment_logsumexp