code: def main():
      # Training settings
      parser = argparse.ArgumentParser(
          description="GNN baselines on pcqm4m with DGL"
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=42,
          help="random seed to use (default: 42)",
      )
      parser.add_argument(
          "--device",
          type=int,
          default=0,
          help="which gpu to use if any (default: 0)",
      )
      parser.add_argument(
          "--gnn",
          type=str,
          default="gin-virtual",
          help="GNN to use, which can be from "
          "[gin, gin-virtual, gcn, gcn-virtual] (default: gin-virtual)",
      )
      parser.add_argument(
          "--graph_pooling",
          type=str,
          default="sum",
          help="graph pooling strategy mean or sum (default: sum)",
      )
      parser.add_argument(
          "--drop_ratio",
          type=float,
          default=0,
          help="dropout ratio (default: 0)",
      )
      parser.add_argument(
          "--num_layers",
          type=int,
          default=5,
          help="number of GNN message passing layers (default: 5)",
      )
      parser.add_argument(
          "--emb_dim",
          type=int,
          default=600,
          help="dimensionality of hidden units in GNNs (default: 600)",
      )
      parser.add_argument(
          "--batch_size",
          type=int,
          default=256,
          help="input batch size for training (default: 256)",
      )
      parser.add_argument(
          "--num_workers",
          type=int,
          default=0,
          help="number of workers (default: 0)",
      )
      parser.add_argument(
          "--checkpoint_dir",
          type=str,
          default="",
          help="directory to save checkpoint",
      )
      parser.add_argument(
          "--save_test_dir",
          type=str,
          default="",
          help="directory to save test submission file",
      )
      args = parser.parse_args()
  
      print(args)
  
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      random.seed(args.seed)
  
      if torch.cuda.is_available():
          torch.cuda.manual_seed(args.seed)
          device = torch.device("cuda:" + str(args.device))
      else:
          device = torch.device("cpu")
  
      ### automatic data loading and splitting
      ### Read in the raw SMILES strings
      smiles_dataset = PCQM4MDataset(root="dataset/", only_smiles=True)
      split_idx = smiles_dataset.get_idx_split()
  
      test_smiles_dataset = [
          smiles_dataset[i] for i in split_idx["test"]
      ]
      onthefly_dataset = OnTheFlyPCQMDataset(test_smiles_dataset)
      test_loader = DataLoader(
          onthefly_dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.num_workers,
          collate_fn=collate_dgl,
      )
  
      ### automatic evaluator.
      evaluator = PCQM4MEvaluator()
  
      shared_params = {
          "num_layers": args.num_layers,
          "emb_dim": args.emb_dim,
          "drop_ratio": args.drop_ratio,
          "graph_pooling": args.graph_pooling,
      }
  
      if args.gnn == "gin":
          model = GNN(
              gnn_type="gin", virtual_node=False, **shared_params
          ).to(device)
      elif args.gnn == "gin-virtual":
          model = GNN(
              gnn_type="gin", virtual_node=True, **shared_params
          ).to(device)
      elif args.gnn == "gcn":
          model = GNN(
              gnn_type="gcn", virtual_node=False, **shared_params
          ).to(device)
      elif args.gnn == "gcn-virtual":
          model = GNN(
              gnn_type="gcn", virtual_node=True, **shared_params
          ).to(device)
      else:
          raise ValueError("Invalid GNN type")
  
      num_params = sum(p.numel() for p in model.parameters())
      print(f"#Params: {num_params}")
  
      checkpoint_path = os.path.join(
          args.checkpoint_dir, "checkpoint.pt"
      )
      if not os.path.exists(checkpoint_path):
          raise RuntimeError(
              f"Checkpoint file not found at {checkpoint_path}"
          )
  
      ## reading in checkpoint
      checkpoint = torch.load(checkpoint_path)
      model.load_state_dict(checkpoint["model_state_dict"])
  
      print("Predicting on test data...")
      y_pred = test(model, device, test_loader)
      print("Saving test submission file...")
      evaluator.save_test_submission(
          {"y_pred": y_pred}, args.save_test_dir
      )
  
file_path: dgl/examples/pytorch/ogb_lsc/PCQM4M/test_inference.py
imports: import argparse
  import os
  import random
  import dgl
  import numpy
  import torch
  from gnn import GNN
  from ogb.lsc import PCQM4MDataset, PCQM4MEvaluator
  from ogb.utils import smiles2graph
  from torch.utils.data import DataLoader
  from tqdm import tqdm
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types