Instead of having the explicit weights W1…W6, maybe you can construct them so that the softmax of them is the weights you want.
def forward(self,x):
    normalized_weights = nn.functional.softmax(self.weights)
    x = nn.functional.linear(x,normalized_weights)