zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix batched add_noise

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-12-08 21:35:16 -08:00
parent 67c68452bb
commit d50ad3ec27
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
flux/flux/sampler.py
View File

@ -50,6 +50,7 @@ class FluxSampler:
if noise is not None if noise is not None
else mx.random.normal(x.shape, dtype=x.dtype, key=key) else mx.random.normal(x.shape, dtype=x.dtype, key=key)
) )
t = t.reshape([-1] + [1] * (x.ndim - 1))
return x * (1 - t) + t * noise return x * (1 - t) + t * noise
def step(self, pred, x_t, t, t_prev): def step(self, pred, x_t, t, t_prev):