Method forward

models/utils.py:199–210

(self, input, i_type)

                nn.init.uniform_(b, -bound, bound)

    def forward(self, input, i_type):
        assert input.dtype == torch.float32 and i_type.dtype == torch.int64
        out = torch.empty(input.shape[0], self.out_features).fill_(np.nan).to(input)
        for t, w in enumerate(self.weight):
            idxs_this_type = (i_type == t).nonzero()[:, 0]
            in_this_type = input[idxs_this_type]
            out_this_type = in_this_type.matmul(w.t())
            if self.bias is not None:
                out_this_type += self.bias[t]
            out[idxs_this_type] = out_this_type

        return out

    # This is a correct but memory-intensive version of forward()
    # def forward(self, input, i_type):