match model forward api

Author: natolambert

src/diffusers/models/unet_rl.py

@@ -130,15 +130,15 @@ def __init__(
             nn.Conv1d(dim, transition_dim, 1),
         )
 
-    def forward(self, sample, timesteps):
+    def forward(self, sample, timestep):
         """
         x : [ batch x horizon x transition ]
         """
         x = sample
 
         x = x.permute(0, 2, 1)
 
-        t = self.time_mlp(timesteps)
+        t = self.time_mlp(timestep)
         h = []
 
         for resnet, resnet2, downsample in self.downs:

tests/test_modeling_utils.py

@@ -585,7 +585,7 @@ def dummy_input(self):
         noise = floats_tensor((batch_size, seq_len, num_features)).to(torch_device)
         time_step = torch.tensor([10] * batch_size).to(torch_device)
 
-        return {"sample": noise, "timesteps": time_step}
+        return {"sample": noise, "timestep": time_step}
 
     @property
     def input_shape(self):