diff --git a/examples/stable_diffusion.py b/examples/stable_diffusion.py
index c484bea53d..f2bc96aabf 100644
--- a/examples/stable_diffusion.py
+++ b/examples/stable_diffusion.py
@@ -250,8 +250,7 @@ class Upsample:
 
 def timestep_embedding(timesteps, dim, max_period=10000):
   half = dim // 2
-  # TODO: remove explicit dtypes after broadcast fix
-  freqs = (-math.log(max_period) * Tensor.arange(half, dtype=dtypes.float32) / half).exp()
+  freqs = (-math.log(max_period) * Tensor.arange(half) / half).exp()
   args = timesteps * freqs
   return Tensor.cat(args.cos(), args.sin()).reshape(1, -1)