github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

hot fix explicitly set arange dtype to float (#2772)

Browse Source
This commit is contained in:
chenyu
2023-12-14 23:14:38 -05:00
committed by GitHub
parent c0f76ed4ea
commit 9afa8009c1
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
examples/stable_diffusion.py
View File

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