Thanks a lot Vahid. That’s a good point. Doing the matrix-matrix multiplication with the shared weights is wasteful. It’s much more efficient to only have a weight vector and then duplicate the outputs, and then add the bias to it. For future reference, the modification (which seems to work) is:

Have only a weight vector and define the bias manually

        self.linear_1 = torch.nn.Linear(1*1*80, 1, bias=False)
        self.linear_1_bias = nn.Parameter(torch.tensor(torch.zeros(num_classes),
                                                       dtype=self.linear_1.weight.dtype))


Then duplicate the outputs over all output units and add the bias vector:

        logits = self.linear_1(out.view(-1, 1*1*80))
        ones = torch.ones(num_classes, dtype=logits.dtype)
        ones = logits
        logits = logits + self.linear_1_bias