Hi, I found out the problem:joy:  the dim in  F.softmax() should be 0 rather than 1. After changing it to 0, the gradient is normal.
Because (bag_embs * self.att_a).mm(self.query_vec) get a Variable with size K * 1 and the softmax should be applied in the first dim.
If F.softmax(..., 1), then the values in alpha are all 1.
BTW, why  the gradient will become 0 when alpha is all 1.