diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
index d790c2c..bcc60b8 100644
--- a/examples/cfd/lid_driven_cavity_2d.py
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -67,7 +67,7 @@ def initialize_fields(self):
         self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.precision_policy, self.backend)
 
     def setup_stepper(self, omega):
-        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions, collision_type='KBC')
 
     def run(self, num_steps, post_process_interval=100):
         for i in range(num_steps):
@@ -100,7 +100,7 @@ def post_process(self, i):
 
         fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude}
 
-        save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity")
+        # save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity")
         save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity")
 
 
@@ -112,7 +112,7 @@ def post_process(self, i):
     precision_policy = PrecisionPolicy.FP32FP32
 
     velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
-    omega = 1.6
+    omega = 1.99
 
     simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
-    simulation.run(num_steps=5000, post_process_interval=1000)
+    simulation.run(num_steps=50000, post_process_interval=1000)
diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py
index ec4f4fe..b841baa 100644
--- a/xlb/operator/collision/kbc.py
+++ b/xlb/operator/collision/kbc.py
@@ -67,7 +67,7 @@ def jax_implementation(
         fneq = f - feq
         if isinstance(self.velocity_set, D2Q9):
             shear = self.decompose_shear_d2q9_jax(fneq)
-            delta_s = shear * rho / 4.0  # TODO: Check this
+            delta_s = shear * rho / 4.0
         elif isinstance(self.velocity_set, D3Q27):
             shear = self.decompose_shear_d3q27_jax(fneq)
             delta_s = shear * rho
@@ -191,16 +191,16 @@ def _construct_warp(self):
         @wp.func
         def decompose_shear_d2q9(fneq: Any):
             pi = self.momentum_flux.warp_functional(fneq)
-            N = pi[0] - pi[1]
+            N = pi[0] - pi[2]
             s = _f_vec()
             s[3] = N
             s[6] = N
             s[2] = -N
             s[1] = -N
-            s[8] = pi[2]
-            s[4] = -pi[2]
-            s[5] = -pi[2]
-            s[7] = pi[2]
+            s[8] = pi[1]
+            s[4] = -pi[1]
+            s[5] = -pi[1]
+            s[7] = pi[1]
             return s
 
         # Construct functional for decomposing shear
@@ -273,9 +273,10 @@ def functional(
             fneq = f - feq
             if wp.static(self.velocity_set.d == 3):
                 shear = decompose_shear_d3q27(fneq)
+                delta_s = shear * rho
             else:
                 shear = decompose_shear_d2q9(fneq)
-            delta_s = shear * rho  # TODO: Check this
+                delta_s = shear * rho / self.compute_dtype(4.0)
 
             # Perform collision
             delta_h = fneq - delta_s