diff --git a/Dockerfile b/Dockerfile
index 97fc0fc..3427650 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,12 +13,13 @@
 # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 # OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-FROM tensorflow/tensorflow:2.7.3-gpu
+FROM tensorflow/tensorflow:2.17.0-gpu
 
 # Need cmake to build the op
 RUN apt-get update && apt-get -y install \
     cmake \
-    libtcmalloc-minimal4
+    libtcmalloc-minimal4 \
+    && rm -rf /var/lib/apt/lists/*
 
 # Need these libraries for training
 RUN pip3 install \
@@ -35,18 +36,18 @@ RUN install -d -m 0777 /.cache/matplotlib
 # Build the tensorflow op and put it in /visualmesh
 RUN mkdir visualmesh
 COPY . visualmesh/
-ENV CXXFLAGS -D_GLIBCXX_USE_CXX11_ABI=0
 RUN mkdir visualmesh/build && cd visualmesh/build \
     && cmake .. \
     -DCMAKE_BUILD_TYPE=Release \
     -DBUILD_EXAMPLES=Off \
+    -DBUILD_OPENCL_ENGINE=Off \
     -DBUILD_TENSORFLOW_OP=On \
     && make
 
-ENV LD_PRELOAD /usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
 
 # Make tensorflow only print out info and above logs
-ENV TF_CPP_MIN_LOG_LEVEL 1
+ENV TF_CPP_MIN_LOG_LEVEL=1
 
 RUN mkdir /workspace
 WORKDIR /workspace
diff --git a/cmake/Modules/FindTensorFlow.cmake b/cmake/Modules/FindTensorFlow.cmake
index 86c14d5..d9d095a 100644
--- a/cmake/Modules/FindTensorFlow.cmake
+++ b/cmake/Modules/FindTensorFlow.cmake
@@ -36,7 +36,7 @@ find_path(
 
 find_library(
     TENSORFLOW_LIBRARIES
-    NAMES tensorflow_framework libtensorflow_framework.so.2 libtensorflow_framework.so.1
+    NAMES tensorflow_framework libtensorflow_framework.so.2 libtensorflow_framework.so.1 libtensorflow_framework.2.dylib
     HINTS ${tf_lib_dir}
     DOC "TensorFlow library")
 
diff --git a/tensorflow/CMakeLists.txt b/tensorflow/CMakeLists.txt
index bbe7b18..9574755 100644
--- a/tensorflow/CMakeLists.txt
+++ b/tensorflow/CMakeLists.txt
@@ -4,6 +4,7 @@ find_package(TensorFlow REQUIRED)
 
 add_library(tf_op SHARED "map.cpp" "unmap.cpp" "lookup.cpp" "difference.cpp" ${hdr})
 target_compile_options(tf_op PRIVATE -march=native -mtune=native)
+target_compile_features(tf_op PRIVATE cxx_std_17)
 set_target_properties(
     tf_op
     PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${PROJECT_SOURCE_DIR}/training/op"
diff --git a/tensorflow/difference.cpp b/tensorflow/difference.cpp
index c4a289a..2d300ce 100644
--- a/tensorflow/difference.cpp
+++ b/tensorflow/difference.cpp
@@ -50,7 +50,7 @@ REGISTER_OP("DifferenceVisualMesh")
   .Output("vectors: T")
   .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
       c->set_output(Outputs::DIFFERENCES, c->input(Args::COORDINATES_A));
-      return tensorflow::Status::OK();
+      return tensorflow::OkStatus();
   });
 
 /**
diff --git a/tensorflow/lookup.cpp b/tensorflow/lookup.cpp
index ba2add0..a1a789b 100644
--- a/tensorflow/lookup.cpp
+++ b/tensorflow/lookup.cpp
@@ -81,7 +81,7 @@ REGISTER_OP("LookupVisualMesh")
       // nx2 vectors on image, n+1xG neighbours (including off screen point), and n global indices
       c->set_output(Outputs::VECTORS, c->MakeShape({kUnknownDim, 3}));
       c->set_output(Outputs::NEIGHBOURS, c->MakeShape({kUnknownDim, kUnknownDim}));
-      return tensorflow::Status::OK();
+      return tensorflow::OkStatus();
   });
 
 /**
diff --git a/tensorflow/map.cpp b/tensorflow/map.cpp
index 840a4e3..5fa39e1 100644
--- a/tensorflow/map.cpp
+++ b/tensorflow/map.cpp
@@ -48,7 +48,7 @@ REGISTER_OP("MapVisualMesh")
   .Output("vectors: T")
   .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
       c->set_output(Outputs::VECTORS, c->Matrix(c->Dim(c->input(Args::COORDINATES), 0), 3));
-      return tensorflow::Status::OK();
+      return tensorflow::OkStatus();
   });
 
 /**
diff --git a/tensorflow/unmap.cpp b/tensorflow/unmap.cpp
index af7ee46..1d1fb07 100644
--- a/tensorflow/unmap.cpp
+++ b/tensorflow/unmap.cpp
@@ -48,7 +48,7 @@ REGISTER_OP("UnmapVisualMesh")
   .Output("coordinates: T")
   .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
       c->set_output(Outputs::COORDINATES, c->Matrix(c->Dim(c->input(Args::VECTORS), 0), 2));
-      return tensorflow::Status::OK();
+      return tensorflow::OkStatus();
   });
 
 /**
diff --git a/training/dataset/visual_mesh_dataset.py b/training/dataset/visual_mesh_dataset.py
index 06f138f..984f26c 100644
--- a/training/dataset/visual_mesh_dataset.py
+++ b/training/dataset/visual_mesh_dataset.py
@@ -170,7 +170,7 @@ def build(self):
         dataset = dataset.filter(lambda args: tf.reduce_all(args["valid"]))
 
         # Perform a ragged batch
-        dataset = dataset.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=self.batch_size))
+        dataset = dataset.ragged_batch(batch_size=self.batch_size)
 
         # Perform actions needed to convert the ragged batches into training examples
         dataset = dataset.map(self._reduce, num_parallel_calls=self.prefetch)
diff --git a/training/metrics/test/bucket.py b/training/metrics/test/bucket.py
index 8efdb4c..604f1c2 100644
--- a/training/metrics/test/bucket.py
+++ b/training/metrics/test/bucket.py
@@ -17,7 +17,7 @@
 
 
 def curve_bucket(x, y):
-    return tf.math.sqrt(tf.add(tf.math.squared_difference(x[:-1], x[1:]), tf.math.squared_difference(y[:-1], y[1:]),))
+    return tf.math.sqrt(tf.add(tf.math.squared_difference(x[:-1], x[1:]), tf.math.squared_difference(y[:-1], y[1:])))
 
 
 def x_bucket(x, y):