How to extract graph node embeddings from a Pytorch-Geometric GAT model?

Dataset Strucute: Temporal directed graph; Nodes have features; Edges don't have features; Nodes are labelled. Using the Elliptic Dataset

Task: Classify nodes/ Predict node labels.

Data Structure: 2 .csv files of nodes and edges.

• For the nodes csv #Rows = #Nodes and #Columns = #Features
• For the edges csv #Rows = #Edges
• Finally both files are converted into a tensor and turned into a Pytorch-geometric Data class.

I want to train various Graph Neural Networks on the data and extract node embeddings from the networks. I know that is possible because the authors of the Elliptic dataset extracted node embeddings from a GCN.

Below is the code for the GAT I am using.

class GAT(torch.nn.Module):"""Graph Attention Network"""  def __init__(self, dim_in, dim_h, dim_out, heads=24):    super().__init__()    self.gat1 = GATv2Conv(dim_in, dim_h, heads=heads)    self.gat2 = GATv2Conv(dim_h*heads, dim_out, heads=1)    self.optimizer = torch.optim.Adam(self.parameters(),                                      lr=0.25,                                      weight_decay=5e-4)  def forward(self, x, edge_index):    h = F.dropout(x, p=0.5, training=self.training)    h = self.gat1(x, edge_index)    h = F.elu(h)    h = F.dropout(h, p=0.5, training=self.training)    h = self.gat2(h, edge_index)    return h, F.log_softmax(h, dim=1)

This function returns a trained model

def train(model, data , epochs = 200):"""Train a GNN model and return the trained model."""    criterion = torch.nn.CrossEntropyLoss()    optimizer = model.optimizer    model = model.to(device)    model.train()    for epoch in range(epochs+1):        # Training        optimizer.zero_grad()        _, out = model(data.x.to(device), data.edge_index.to(device))        loss = criterion(out[data.train_mask].to(device), data.y[data.train_mask].to(device))        loss.backward()        optimizer.step()        # Print metrics every 10 epochs        if(epoch % 10 == 0):            print(f'Epoch {epoch:>3} | Train Loss: {loss:.3f}')    return model

What modifications do I need to make to the code to extract the node embeddings?