Thank you so much.
Do you know how to translate this piece of code to PyTorch?

  
      github.com
  
  
    facebookresearch/CommNet/blob/d7450fc3cd336119ea5bf970a0db1added65d2ca/linear_multi.lua 1
-- Copyright (c) 2016-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is licensed under the BSD-style license found in the
-- LICENSE file in the root directory of this source tree. An additional grant
-- of patent rights can be found in the PATENTS file in the same directory.

-- linear layer with multple copies of parameters
-- which parameters to use is set by model_ids
function linear_multi(sz_in, sz_out, model_ids, input)
	if g_opts.nmodels == 1 then
		return nn.Linear(sz_in, sz_out)(input)
	end
    local weight_lut = nn.LookupTable(g_opts.nmodels, sz_in * sz_out)(model_ids)
    weight_lut.data.module.updateGradInput = function(self, input) return input end
    local bias_lut = nn.LookupTable(g_opts.nmodels, sz_out)(model_ids)
    bias_lut.data.module.updateGradInput = function(self, input) return input end
    local weight_view = nn.View(sz_out, sz_in):setNumInputDims(1)(weight_lut)
    input = nn.View(-1, 1):setNumInputDims(1)(input)
    local out = nn.MM(false, false)({weight_view, input})


  This file has been truncated. show original

  
  
    
    
  
  


Here is what I have done so far:


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)

    weight_lut = nn.Embedding(nmodels, sz_in * sz_out)(model_ids)
    bias_lut = nn.Embedding(nmodels, sz_out)(model_ids)
    weight_view = weight_lut.view(sz_out, sz_in)

    a, b, c = input.size()
    input = input.view(a*b, c)

    out = torch.mm(weight_view, input)

    a, b, c, d = out.size()
    out = out.view(-1)

    out = out.add(bias_lut)
    out.weight_lut = weight_lut
    out.bias_lut = bias_lut



class TestModule(nn.Module):
    def __init__(self, nmodels):
        # g_opts = {nmodels = 2}
        # lm = linear_multi(10, 20, ids, input)
        # w = lm.weight_lut.data.module
        # b = lm.bias_lut.data.module
        # w.weight[1]:copy(l1.weight:view(-1))
        # w.weight[2]:copy(l2.weight:view(-1))
        # b.weight[1]:copy(l1.bias)
        # b.weight[2]:copy(l2.bias)
        # model = nn.gModule({input, ids}, {lm})
        self.nmodels = nmodels

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

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(2)
    model.forward(torch.LongTensor([1,2,1]), x)