code: def train_hetero(
      model: Any,
      loader: NeighborLoader,
      optimizer: torch.optim.Adam,
      device: torch.device,
  ) -> torch.Tensor:
      for batch in loader:
          optimizer.zero_grad()
          batch = batch.to(device)
          out = model(batch.x_dict, batch.edge_index_dict)
          batch_size = batch["paper"].batch_size
          out = out["paper"][:batch_size]
          target = batch["paper"].y[:batch_size]
          loss = F.cross_entropy(out, target)
          loss.backward()
          optimizer.step()
  
      return loss
  
file_path: pytorch_geometric/benchmark/multi_gpu/training/common.py
imports: import argparse
  import ast
  from time import perf_counter
  from typing import Any, Callable, Tuple, Union
  import torch
  import torch.distributed
  import torch.nn.functional
  from torch.nn.parallel import DistributedDataParallel
  from benchmark.utils import get_model, get_split_masks, test
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import PNAConv
label_desc: Use image node_id 2 for calling a global function with example usage: train_hetero(model, loader, optimizer, device) and returns: loss