Hello, I am experiencing a similar issue, however I do not fully understand the behavior, so I would be grateful for some insight.
When I run the code with these 2 lines of code, it causes the mentioned In-place operation error
attention = nn.functional.softmax(torch.matmul(key.transpose(1, 2), query) / math.sqrt(num_of_channels), dim=1)
attention[:, 0, 0] = torch.ones(attention.shape[0], requires_grad = True)

Error:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

However, if I run this code like this:
attention = nn.functional.softmax(torch.matmul(key.transpose(1, 2), query), dim=1) / math.sqrt(num_of_channels)
attention[:, 0, 0] = torch.ones(attention.shape[0], requires_grad = True)

it works perfectly fine.
What is the reason for such behavior? The resulting tensor has the attribute requires_grad set to True in both cases, only the backward function differs. Actually, what is fun is that I can just use a dummy division by 1.0 to avoid getting an error, but it does not seem elegant at all:
 attention = nn.functional.softmax(torch.matmul(key.transpose(1, 2), query) / math.sqrt(num_of_channels), dim=1)/1.0
attention[:, 0, 0] = torch.ones(attention.shape[0], requires_grad = True)