code: def train():
      model.train()
  
      total_examples = total_loss = 0
      for batch in tqdm(train_loader):
          batch = batch.to(device)
          optimizer.zero_grad()
  
          out = model(
              batch.x_dict,
              batch.adj_t_dict
              if args.use_sparse_tensor
              else batch.edge_index_dict,
              num_sampled_nodes_dict=batch.num_sampled_nodes_dict,
              num_sampled_edges_dict=batch.num_sampled_edges_dict,
          )
  
          batch_size = batch["paper"].batch_size
          loss = F.cross_entropy(
              out[:batch_size], batch["paper"].y[:batch_size]
          )
          loss.backward()
          optimizer.step()
  
          total_examples += batch_size
          total_loss += float(loss) * batch_size
  
      return total_loss / total_examples
  
file_path: pytorch_geometric/examples/hetero/hierarchical_sage.py
imports: import argparse
  import torch
  import torch.nn.functional
  from tqdm import tqdm
  import torch_geometric.transforms
  from torch_geometric.datasets import OGB_MAG
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import HeteroConv, Linear, SAGEConv
  from torch_geometric.utils import trim_to_layer
label_desc: Use image node_id 1 for calling a global function with example usage: train() and returns: unknown