code: def train():
      memory.train()
      gnn.train()
      link_pred.train()
  
      memory.reset_state()  # Start with a fresh memory.
      neighbor_loader.reset_state()  # Start with an empty graph.
  
      total_loss = 0
      for batch in train_loader:
          optimizer.zero_grad()
          batch = batch.to(device)
  
          n_id, edge_index, e_id = neighbor_loader(batch.n_id)
          assoc[n_id] = torch.arange(n_id.size(0), device=device)
  
          # Get updated memory of all nodes involved in the computation.
          z, last_update = memory(n_id)
          z = gnn(
              z,
              last_update,
              edge_index,
              data.t[e_id].to(device),
              data.msg[e_id].to(device),
          )
          pos_out = link_pred(z[assoc[batch.src]], z[assoc[batch.dst]])
          neg_out = link_pred(
              z[assoc[batch.src]], z[assoc[batch.neg_dst]]
          )
  
          loss = criterion(pos_out, torch.ones_like(pos_out))
          loss += criterion(neg_out, torch.zeros_like(neg_out))
  
          # Update memory and neighbor loader with ground-truth state.
          memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
          neighbor_loader.insert(batch.src, batch.dst)
  
          loss.backward()
          optimizer.step()
          memory.detach()
          total_loss += float(loss) * batch.num_events
  
      return total_loss / train_data.num_events
  
file_path: pytorch_geometric/examples/tgn.py
imports: import os.path
  import torch
  from sklearn.metrics import average_precision_score, roc_auc_score
  from torch.nn import Linear
  from torch_geometric.datasets import JODIEDataset
  from torch_geometric.loader import TemporalDataLoader
  from torch_geometric.nn import TGNMemory, TransformerConv
  from torch_geometric.nn.models.tgn import IdentityMessage, LastAggregator, LastNeighborLoader
label_desc: Use image node_id 1 for calling a global function with example usage: train() and returns: unknown