Fix flux training with batch size (#1135)

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

flux/flux/sampler.py

@@ -50,6 +50,7 @@ def add_noise(self, x, t, noise=None, key=None):
             if noise is not None
             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
 
     def step(self, pred, x_t, t, t_prev):