Thank you very much.
After 2 hours, this seems to be working:

from torch import nn
from torch.autograd import Variable
import torch

def linear_multi(nmodels, sz_in, sz_out, model_ids, input):
    if nmodels == 1:
        return nn.Linear(sz_in, sz_out)(input)

    # XXX: potential bug - updateGradInput is overridden,
    # possible use of `register_backward_hook`

    weight_lut = nn.Embedding(nmodels, sz_in * sz_out)(model_ids) # 1x3x200
    bias_lut = nn.Embedding(nmodels, sz_out)(model_ids) # 1x3x20

    weight_view = weight_lut.view(-1, sz_in, sz_out) # 3 x 10 x 20
    bias_view = bias_lut.view(-1, sz_out) # 3x20

    a, b = input.size()
    input = input.view(a, 1, b) # 3x1x10

    out = torch.matmul(input, weight_view)# 3x1x20

    a, b, c = out.size()
    out = out.view(a, c) #3x20
    out = out.add(bias_view) # 3x20
    return out



class TestModule(nn.Module):
    """
        g_opts = {nmodels = 2}
        lm = linear_multi(10, 20, ids, input)
        model = nn.gModule({input, ids}, {lm})
    """
    def __init__(self, nmodels):
        self.nmodels = nmodels

    def forward(self, input, ids):
        lm = linear_multi(self.nmodels, 10, 20, ids, input)
        return lm

if __name__ == "__main__":
    l1 = nn.Linear(10, 20)
    l2 = nn.Linear(10, 20)
    x = Variable(torch.rand(3, 10))
    y1 = l1.forward(x)
    y2 = l2.forward(x)

    model = TestModule(3) # in lua: 0, 1, 2
    y = model.forward(x, Variable(torch.LongTensor([[1,2,1]])))
    assert y.size()[0] == 3
    assert y.size()[1] == 20
    assert y.dim() == 2
    
    # Note: in the original test, the weight of l1, l2 is copied to the
    # weight of linear_multi. Then test the matmul results are the same