diff --git a/README.md b/README.md
index e2c4fa9..8d662b5 100644
--- a/README.md
+++ b/README.md
@@ -144,7 +144,7 @@ Please refer to https://github.com/google/jax for the latest installation docume
 
 Install dependencies:
 ```bash
-pip install jmp pyvista numpy matplotlib Rtree trimesh jmp
+pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor
 ```
 
 Run an example:
@@ -152,5 +152,5 @@ Run an example:
 git clone https://github.com/Autodesk/XLB
 cd XLB
 export PYTHONPATH=.
-python3 examples/cavity2d.py
+python3 examples/CFD/cavity2d.py
 ```
diff --git a/src/utils.py b/src/utils.py
index 61c11fb..1720ca6 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -123,7 +123,7 @@ def save_fields_vtk(timestep, fields, output_dir='.', prefix='fields'):
     if value.ndim == 2:
         dimensions = dimensions + (1,)
 
-    grid = pv.UniformGrid(dimensions=dimensions)
+    grid = pv.ImageData(dimensions=dimensions)
 
     # Add the fields to the grid
     for key, value in fields.items():
@@ -157,7 +157,7 @@ def live_volume_randering(timestep, field):
     if field.ndim != 3:
         raise ValueError("The input field must be 3D!")
     dimensions = field.shape
-    grid = pv.UniformGrid(dimensions=dimensions)
+    grid = pv.ImageData(dimensions=dimensions)
 
     # Add the field to the grid
     grid['field'] = field.flatten(order='F')
@@ -207,7 +207,7 @@ def save_BCs_vtk(timestep, BCs, gridInfo,  output_dir='.'):
         gridDimensions = (gridInfo['nx'] + 1, gridInfo['ny'] + 1, gridInfo['nz'] + 1)
         fieldDimensions = (gridInfo['nx'], gridInfo['ny'], gridInfo['nz'])
 
-    grid = pv.UniformGrid(dimensions=gridDimensions)
+    grid = pv.ImageData(dimensions=gridDimensions)
 
     # Dictionary to keep track of encountered BC names
     bcNamesCount = {}