Method __init__

tests/test_parametrize.py:212–215

(self, device, dtype)

    class SimpleMoE(nn.Module):
        def __init__(self, device, dtype):
            super().__init__()
            # Expert weights: [num_experts, input_dim, output_dim]
            self.expert_weights = nn.Parameter(torch.randn(4, 32, 64, dtype=dtype, device=device))

        def forward(self, x, expert_idx=0):
            # Select and use specific expert weight matrix