I think you can combine softmax with your parameters.
Maybe you should define your own module, and do following:
def forward(self, x):
    w_normalized = nn.functional.softmax(self.weight)
    output = nn.functional.linear(x, w_normalized)

    return output

Then, you can train the model that satisfies your constraint.