帮我生产一段GNN代码

好的，以下是一个简单的GNN代码示例，用于节点分类任务：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GNN, self).__init__()
        self.conv1 = GCNConv(input_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, output_dim)
        
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

这个模型使用了两层GCN（Graph Convolutional Network）层，输入维度为`input_dim`，隐藏层维度为`hidden_dim`，输出维度为`output_dim`。在`forward`函数中，我们通过两个GCN层和ReLU激活函数来对节点特征进行编码。最后通过log_softmax函数来输出预测结果。