Hi @albanD  and thanks for for fast answers! I actually just want to take the maximum element from matrix and ignore everything else. But i still need the gradients to flow through W because I need it to learn how to choose the maximum.
outputs = softmax(A.W.B)
value, idx = torch.max(outputs)
mask = torch.ones_like(outputs)
mask[outputs<value] = 0
mask[outputs>=value] = 1
outputs = outputs * mask
weighted_sum = torch.sum(outputs * matrix)

This would basically keep only the maximum element. This would be fine I guess? We are just changing the output tensor of a specific layer.