Good idea!
Unfortunately the values of a don’t change. Here is the relevant part of my code (I’m using PyTorch 0.1.12):
self.weight_vector = Variable(torch.FloatTensor(1, 4), requires_grad=True)
def forward(self, v): # v.size() = (4,36,256)
    norm_weight_vec = F.softmax(self.weight_vector)
    v = v * (norm_weight_vec.transpose(0, 1).unsqueeze(2)).expand_as(x)
    v = torch.sum(x, dim=0) # result size: (1,36,256)
    ....