replace unsued arg

Author: younesbelkada

src/diffusers/models/unet_2d_condition_flax.py

@@ -192,7 +192,7 @@ def setup(self):
     def __call__(
         self,
         sample,
-        timestep,
+        timesteps,
         encoder_hidden_states,
         return_dict: bool = True,
         train: bool = False,
@@ -214,11 +214,10 @@ def __call__(
             When returning a tuple, the first element is the sample tensor.
         """
         # 1. time
-        timesteps = timestep
         if not isinstance(timesteps, jnp.ndarray):
             timesteps = jnp.array([timesteps], dtype=jnp.int32)
         elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0:
-            timesteps = timesteps.to(dtype=jnp.float32)
+            timesteps = timesteps.astype(dtype=jnp.float32)
             timesteps = timesteps[None]
 
         t_emb = self.time_proj(timesteps)