I see. One way is that, you can consider self.weight as logits and apply softmax in forward to get the adjacency matrix. (not sure why do we need relu), then continue with GNN operation.
But you do not need to assign it back to self.weight, as the logits will be learned when you loss.backward() and optimizer.step()