帮我用pytorch写一个图注意力网络的代码

好的，以下是一个简单的图注意力网络的 PyTorch 代码：

import torch
import torch.nn as nn

class GraphAttention(nn.Module):
    def __init__(self, in_features, out_features):
        super(GraphAttention, self).__init__()
        self.W = nn.Parameter(torch.Tensor(in_features, out_features))
        self.a = nn.Parameter(torch.Tensor(2*out_features, 1))

        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        nn.init.xavier_uniform_(self.a.data, gain=1.414)

    def forward(self, inputs):
        h = torch.mm(inputs, self.W)
        attention = torch.cat([h, h], dim=-1) @ self.a
        attention = torch.softmax(attention, dim=1)
        output = torch.transpose(inputs, 0, 1) @ attention
        return output

class GraphAttentionNetwork(nn.Module):
    def __init__(self, in_features, out_features, num_heads):
        super(GraphAttentionNetwork, self).__init__()
        self.attention_layers = nn.ModuleList([GraphAttention(in_features, out_features) for _ in range(num_heads)])
        self.output_layer = nn.Linear(num_heads * out_features, out_features)

    def forward(self, inputs):
        outputs = [attention(inputs) for attention in self.attention_layers]
        output = torch.cat(outputs, dim=-1)
        output = self.output_layer(output)
        return output

该代码实现了一个简单的图注意力网络，其中包含多个注意力头，每个头都会计算一个注意力系数，然后将它们合并成一个输出。与 TensorFlow 版本的代码类似，这里也使用了 PyTorch 的 nn.Module 和 nn.Parameter 来实现可训练的参数。