code: def train():
      model.train()
  
      model.mem1.k.requires_grad = False
      model.mem2.k.requires_grad = False
      for data in train_loader:
          optimizer.zero_grad()
          data = data.to(device)
          out = model(data.x, data.edge_index, data.batch)[0]
          loss = F.nll_loss(out, data.y)
          loss.backward()
          optimizer.step()
  
      kl_loss = 0.0
      model.mem1.k.requires_grad = True
      model.mem2.k.requires_grad = True
      optimizer.zero_grad()
      for data in train_loader:
          data = data.to(device)
          kl_loss += model(data.x, data.edge_index, data.batch)[1]
      kl_loss /= len(train_loader.dataset)
      kl_loss.backward()
      optimizer.step()
  
file_path: pytorch_geometric/examples/mem_pool.py
imports: import os.path
  import time
  import torch
  import torch.nn.functional
  from torch.nn import BatchNorm1d, LeakyReLU, Linear
  from torch_geometric.datasets import TUDataset
  from torch_geometric.loader import DataLoader
  from torch_geometric.nn import DeepGCNLayer, GATConv, MemPooling
label_desc: Use image node_id 1 for calling a global function with example usage: train() without return types