diff --git a/mlx/fast.cpp b/mlx/fast.cpp
index 803ec57177..2afb834166 100644
--- a/mlx/fast.cpp
+++ b/mlx/fast.cpp
@@ -352,7 +352,7 @@ array layer_norm(
     auto xc = subtract(x, mu, s);
     auto v = mean(square(xc, s), /* axis= */ -1, /* keepdims= */ true, s);
 
-    x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s));
+    x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s), s);
     x = astype(x, out_type, s);
 
     // If the LN is affine then transform x according to the weight and bias
@@ -945,8 +945,9 @@ array scaled_dot_product_attention(
       array& mask = inputs[3];
       mask = where(
           mask,
-          full_like(mask, 0, final_type, s),
-          full_like(mask, -inf, final_type, s));
+          full_like(mask, 0, final_type, stream),
+          full_like(mask, -inf, final_type, stream),
+          stream);
     }
     Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)};
     auto primitive = std::make_shared<ScaledDotProductAttention>(