code: def train_hetero(
      model,
      loader,
      optimizer,
      device,
      progress_bar=True,
      desc="",
      trim=False,
  ):
      if trim:
          warnings.warn(
              "Trimming not yet implemented for heterogeneous graphs"
          )
  
      if progress_bar:
          loader = tqdm(loader, desc=desc)
      for batch in loader:
          optimizer.zero_grad()
          batch = batch.to(device)
          if "adj_t" in batch:
              edge_index_dict = batch.adj_t_dict
          else:
              edge_index_dict = batch.edge_index_dict
          out = model(batch.x_dict, 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()
  
file_path: pytorch_geometric/benchmark/training/training_benchmark.py
imports: import argparse
  import ast
  import warnings
  from collections import defaultdict
  from contextlib import nullcontext
  import torch
  import torch.nn.functional
  from tqdm import tqdm
  from benchmark.utils import emit_itt, get_dataset, get_model, get_split_masks, save_benchmark_data, test, write_to_csv
  from torch_geometric import compile
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import PNAConv
  from torch_geometric.profile import rename_profile_file, timeit, torch_profile, xpu_profile
label_desc: Use image node_id 2 for calling a global function with example usage: train_hetero(model, loader, optimizer, device, progress_bar, desc, trim) without return types