From 30fc1ad415d0914b0a8860ccb3cf58f5aa02fa4a Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Sun, 21 Apr 2024 00:31:24 -0400
Subject: [PATCH] remove TODO: remove explicit dtypes after broadcast fix in
 stable_diffusion (#4241)

this is done
---
 examples/stable_diffusion.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

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)