(self, dim: int, num_experts: int, top_k: int = 2)
| 91 | """ |
| 92 | |
| 93 | def __init__(self, dim: int, num_experts: int, top_k: int = 2): |
| 94 | super(BitMoE, self).__init__() |
| 95 | self.router = NoisyTopkRouter(dim, num_experts, top_k) |
| 96 | self.experts = nn.ModuleList([Expert(dim) for _ in range(num_experts)]) |
| 97 | self.top_k = top_k |
| 98 | |
| 99 | def forward(self, x): |
| 100 | gating_output, indices = self.router(x) |
nothing calls this directly
no test coverage detected