How to constrict each column of weight matrix sum to 1

2

Did you try this?
https://discuss.pytorch.org/t/contraining-weights-to-sum-to-1/20609/2