Upload moe.py

moe.py

@@ -86,13 +86,14 @@ class MaskedMoE(MoE):
 
 
     def forward(self, inputs: torch.Tensor, mask: torch.Tensor):
+        seq_len = inputs.shape[1]
         inputs_squashed = inputs.view(-1, inputs.shape[-1])
         router_logits = self.router(inputs_squashed)
         mask = torch.cat(
             (mask, torch.ones((mask.shape[0], 1), device=mask.device)),
             dim=1
         )
-        mask = mask.repeat_interleave(self._sequence_length, dim=0)
+        mask = mask.repeat_interleave(seq_len, dim=0)
         router_logits = router_logits*mask
 
         # note that selected experts will be the same for all orders: