给出图神经网络关于时空图的模型代码，用pytorch实现

以下是一个基于PyTorch的图神经网络模型代码，用于处理时空图数据：

import torch
import torch.nn as nn
import torch.nn.functional as F

class STGNN(nn.Module):
    def __init__(self, num_nodes, in_dim, hidden_dim, out_dim, dropout_prob):
        super(STGNN, self).__init__()
        self.num_nodes = num_nodes
        self.in_dim = in_dim
        self.hidden_dim = hidden_dim
        self.out_dim = out_dim
        self.dropout_prob = dropout_prob
        
        # Spatial Temporal Convolutional Layers
        self.st_conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(3,3), padding=1)
        self.st_conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3,3), padding=1)
        
        # Graph Convolutional Layers
        self.gc1 = GCNConv(in_channels=in_dim, out_channels=hidden_dim)
        self.gc2 = GCNConv(in_channels=hidden_dim, out_channels=out_dim)
        
        # Dropout 
        self.dropout = nn.Dropout(p=dropout_prob)
        
    def forward(self, x, adj):
        # Reshape the input data to [batch_size, 1, num_nodes, num_time_steps]
        x = x.reshape(-1, 1, self.num_nodes, -1)
        
        # Spatial Temporal Convolutional Layers
        x = F.relu(self.st_conv1(x))
        x = F.relu(self.st_conv2(x))
        
        # Reshape the output data from [batch_size, num_filters, num_nodes, num_time_steps] to [batch_size, num_nodes, num_time_steps, num_filters]
        x = x.permute(0, 2, 3, 1)
        x = x.reshape(-1, x.shape[1], x.shape[2]*x.shape[3])
        
        # Graph Convolutional Layers
        x = self.gc1(x, adj)
        x = F.relu(x)
        x = self.dropout(x)
        
        x = self.gc2(x, adj)
        x = F.relu(x)
        x = self.dropout(x)
        
        # Reshape the output data from [batch_size*num_nodes, out_dim] to [batch_size, num_nodes*out_dim]
        x = x.reshape(-1, self.num_nodes*self.out_dim)
        
        return x

在这个模型中，我们首先使用两个卷积层来处理时空图的输入数据。然后，我们将处理后的数据送入两个图卷积层中进行进一步的处理。

在这个模型中，我们使用了 PyTorch Geometric 库中的 GCNConv 模块来实现图卷积层。这个模块可以直接接受邻接矩阵作为输入，因此我们可以直接将邻接矩阵作为参数传递给 GCNConv 模块。

最后，我们对模型的输出进行了一些格式转换，以便于进行后续的计算和分析。