diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
index e1d5f64..e7fab29 100644
--- a/examples/CFD/cylinder2d.py
+++ b/examples/CFD/cylinder2d.py
@@ -50,8 +50,9 @@ def set_boundary_conditions(self):
 
         # Outflow BC
         outlet = self.boundingBoxIndices['right']
-        rho_outlet = np.ones(outlet.shape[0], dtype=self.precisionPolicy.compute_dtype)
+        rho_outlet = np.ones((outlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
         self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
+        # self.BCs.append(ZouHe(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet))
 
         # Inlet BC
         inlet = self.boundingBoxIndices['left']
diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
index cd98c91..a2d0087 100644
--- a/src/boundary_conditions.py
+++ b/src/boundary_conditions.py
@@ -702,11 +702,11 @@ def calculate_vel(self, fpop, rho):
         """
         Calculate velocity based on the prescribed pressure/density (Zou/He BC)
         """
-        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1) +
-                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1))
+        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
+                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
 
         # Return the above unormal as a normal vector which sets the tangential velocities to zero
-        vel = unormal[:, None] * self.normals
+        vel = unormal * self.normals
         return vel
 
     @partial(jit, static_argnums=(0,), inline=True)