LookupTable.py 文件源码

def updateOutput(self, input):
        self.renorm(input)
        input = self._makeInputContiguous(input)
        if input.dim() == 1:
           torch.index_select(self.output, self.weight, 0, input)
        elif input.dim() == 2:
           torch.index_select(self.output, self.weight, 0, input.view(-1))
           self.output = self.output.view(input.size(0), input.size(1), self.weight.size(1))
        else:
           raise RuntimeError("input must be a vector or matrix")

        return self.output