Update model.py

model.py

@@ -355,7 +355,7 @@ class StripedHyena(nn.Module):
         self.gradient_checkpointing = False
         self._gradient_checkpointing_func = None
 
-    def forward(self, inputs_embeds, inference_params_dict=None, padding_mask=None):
+    def forward(self, x, inference_params_dict=None, padding_mask=None):
         L = x.shape[1]
         x = self.embedding_layer.embed(x)
         if inference_params_dict is not None:
@@ -370,7 +370,7 @@ class StripedHyena(nn.Module):
         x = self.unembed.unembed(x)
         return x, inference_params_dict_out
 
-    def stateful_forward(self, inputs_embeds, inference_params_dict=None):
+    def stateful_forward(self, x, inference_params_dict=None):
         for block_idx, block in enumerate(self.blocks):
             block_name = "mha" if block_idx in self.config.attn_layer_idxs else "hyena"
             inference_params = inference_params_dict[block_name]
@@ -378,7 +378,7 @@ class StripedHyena(nn.Module):
 
         return x, inference_params_dict
 
-    def stateless_forward(self, inputs_embeds, padding_mask=None):
+    def stateless_forward(self, x, padding_mask=None):
         if type(padding_mask) == torch.Tensor:
             x = x * padding_mask[..., None]