Thanks for the code.
The issue is created by reusing K and assigning some results to it in:
K[torch.broadcast_tensors(w3 < self.att_th, K)[0]] = -1.0

Since K doesn’t require gradients, the return value of HardSelfAttentionLayer will also not require gradients.
I’m not familiar with your approach, but the index operation won’t be differentiable.