diff --git a/xlb/operator/boundary_condition/bc_regularized.py b/xlb/operator/boundary_condition/bc_regularized.py
index af4c783..5fd71f6 100644
--- a/xlb/operator/boundary_condition/bc_regularized.py
+++ b/xlb/operator/boundary_condition/bc_regularized.py
@@ -162,9 +162,14 @@ def _get_fsum(
         def get_normal_vectors(
             missing_mask: Any,
         ):
-            for l in range(_q):
-                if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
-                    return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
+            if wp.static(_d == 3):
+                for l in range(_q):
+                    if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
+                        return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
+            else:
+                for l in range(_q):
+                    if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1:
+                        return -_u_vec(_c_float[0, l], _c_float[1, l])
 
         @wp.func
         def bounceback_nonequilibrium(
diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py
index 5c5c29e..ec4f4fe 100644
--- a/xlb/operator/collision/kbc.py
+++ b/xlb/operator/collision/kbc.py
@@ -271,7 +271,10 @@ def functional(
         ):
             # Compute shear and delta_s
             fneq = f - feq
-            shear = decompose_shear_d3q27(fneq)
+            if wp.static(self.velocity_set.d == 3):
+                shear = decompose_shear_d3q27(fneq)
+            else:
+                shear = decompose_shear_d2q9(fneq)
             delta_s = shear * rho  # TODO: Check this
 
             # Perform collision