def _top_k(self, scores):
        if self.args.moe_top_k == 1:
            return scores.max(dim=-1) # <-- causes weight shape to become [S]
        return torch.topk(scores, self.args.moe_top_k, dim=-1) # <-- shape is normally [S,K]

        expert_weights, expert_indices = self._top_k(scores)
        if self.args.moe_normalize_expert_weights:
            # this function expects dim=-1 to only contain a single token's weights
            expert_weights = expert_weights / torch.norm(
                expert_weights, p=self.args.moe_normalize_expert_weights,dim=-1, keepdim=True)