def METHOD_NAME(self, input, bs, prototypes, outer_batch=False, dispatchV2=False):
    # sm -> se -> szf
    logits = self.wg(input)
    logits = logits.reshape(input.shape[0], prototypes, -1)
    if outer_batch:
        raise NotImplementedError
    else:
        return topk_gating(
            logits,
            k=self.k,
            capacity=self.capacity,
            need_l_aux=self.need_l_aux,
            dispatchV2=dispatchV2,
        )