code: def main():
      # Training settings
      parser = argparse.ArgumentParser(
          description="GNN baselines on pcqm4m with Pytorch Geometrics"
      )
      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 gin, gin-virtual, or gcn, or 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(42)
      torch.manual_seed(42)
      torch.cuda.manual_seed(42)
      random.seed(42)
  
      device = (
          torch.device("cuda:" + str(args.device))
          if torch.cuda.is_available()
          else torch.device("cpu")
      )
  
      ### automatic dataloading 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-whole"]
      ]
      onthefly_dataset = OnTheFlyPCQMDataset(test_smiles_dataset)
      test_loader = DataLoader(
          onthefly_dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.num_workers,
      )
  
      ### automatic evaluator. takes dataset name as input
      evaluator = PCQM4MEvaluator()
  
      if args.checkpoint_dir is not "":
          os.makedirs(args.checkpoint_dir, exist_ok=True)
  
      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, mode="test-whole"
      )
  
imports: import torch
  from torch_geometric.data import DataLoader
  import torch.optim
  import torch.nn.functional
  from torch.utils.tensorboard import SummaryWriter
  from torch.optim.lr_scheduler import StepLR
  from gnn import GNN
  import os
  from tqdm import tqdm
  import argparse
  import time
  import numpy
  import random
  from ogb.lsc import PCQM4MDataset, PCQM4MEvaluator
  from ogb.utils import smiles2graph
  from torch_geometric.data import Data
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types