(self, x, expert_idx=0)
| 215 | self.expert_weights = nn.Parameter(torch.randn(4, 32, 64, dtype=dtype, device=device)) |
| 216 | |
| 217 | def forward(self, x, expert_idx=0): |
| 218 | # Select and use specific expert weight matrix |
| 219 | expert_weight = self.expert_weights[expert_idx] # Shape: [input_dim, output_dim] |
| 220 | return torch.matmul(x, expert_weight) |
| 221 | |
| 222 | module = SimpleMoE(device=device, dtype=dtype) |
| 223 | x = torch.randn(8, 32, dtype=dtype, device=device) |
nothing calls this directly
no outgoing calls
no test coverage detected