def METHOD_NAME(self, data) -> torch.Tensor:
    num_neighbors = 2  # typical number of neighbors
    num_nodes = 4  # typical number of nodes
    edge_src, edge_dst = radius_graph(x=data.pos, r=1.1, batch=data.batch)  # tensors of indices representing the graph
    edge_vec = data.pos[edge_src] - data.pos[edge_dst]
    edge_sh = o3.spherical_harmonics(
        l=self.irreps_sh,
        x=edge_vec,
        normalize=False,  # here we don't normalize otherwise it would not be a polynomial
        normalization="component",
    )
    # For each node, the initial features are the sum of the spherical harmonics of the neighbors
    node_features = scatter(edge_sh, edge_dst, dim=0).div(num_neighbors**0.5)
    # For each edge, tensor product the features on the source node with the spherical harmonics
    edge_features = self.tp1(node_features[edge_src], edge_sh)
    node_features = scatter(edge_features, edge_dst, dim=0).div(num_neighbors**0.5)
    edge_features = self.tp2(node_features[edge_src], edge_sh)
    node_features = scatter(edge_features, edge_dst, dim=0).div(num_neighbors**0.5)
    # For each graph, all the node's features are summed
    return scatter(node_features, data.batch, dim=0).div(num_nodes**0.5)