Am I doing the right thing by classifying heterogeneous graphs?Maybe there's another way?

[Gpy23]

The data is constructed in this way
data = HeteroData()
This is how the network structure is defined

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1,-1), hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)
        self.lin = Linear(hidden_channels, out_channels)
    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = global_mean_pool(x, batch)
        x = self.lin(x)
        return x

Because we have to do graph classification, so a global pooling layer is added
x = global_mean_pool(x, batch)
In doing this stepmodel = to_hetero(model, data.metadata(), aggr='sum')
there will be such errorsx = global_mean_pool(x, batch)
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow