fix: local

modeling_hf_nomic_bert.py

@@ -1244,18 +1244,10 @@ class NomicMoELayer(nn.Module):
 
     def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
         batch_size, seq_len, hidden_dim = x.shape
-        if attention_mask is not None:
-            valid_indices = attention_mask.bool().view(-1)
-            x_valid = x.view(-1, hidden_dim)[valid_indices]
 
         weights, top_weights, top_experts = self.router(x)
         out = self.experts(x, weights, top_weights, top_experts)
 
-        if attention_mask is not None:
-            full_out = torch.zeros(batch_size * seq_len, hidden_dim, dtype=out.dtype, device=out.device)
-            full_out[valid_indices] = out
-            out = full_out.view(batch_size, seq_len, hidden_dim)
-
         return out