I want to add a pooling layer on my model. I want to preserve the num of node as I m doing node prediction without batching.

import torch
from torch_geometric.nn import global_add_pool
# Assume your input tensor is x with shape (num_nodes, num_node_features)
x = torch.randn(10, 10)  # Example tensor with (num_nodes, num_node_features) = (10, 10)
# Apply global_add_pool
x_pooled = global_add_pool(x, batch=None)

x_pooled (1,10) But I want a pooling method to give me (10,x)