I find that add a line before
message_weight[:, 0, :, :] = 0.5

such that
message_weight = message_weight.clone()
message_weight[:, 0, :, :] = 0.5

autograd works.
I am wondering why clone is not necessary in the code you supplied.