small change

- use `jnp.expand_dims` for converting `timesteps` in case it is a 0-dimensional array

Author: younesbelkada

src/diffusers/models/unet_2d_condition_flax.py

@@ -218,7 +218,7 @@ def __call__(
             timesteps = jnp.array([timesteps], dtype=jnp.int32)
         elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0:
             timesteps = timesteps.astype(dtype=jnp.float32)
-            timesteps = timesteps[None]
+            timesteps = jnp.expand_dims(timesteps, 0)
 
         t_emb = self.time_proj(timesteps)
         t_emb = self.time_embedding(t_emb)