I would remove nn.Linear altogether, and have something like
w2 = self.weight_b if stage2 else self.weight_a
x_b = x.matmul(w2.t())

another alternative that needs no flag passing:
w2 = self.weight_a + self.weight_b

init weight_b to zero and don’t train it in phase 1, detach or otherwise freeze weight_a in phase 2