code: def main_mlp():
      # Training settings
      parser = argparse.ArgumentParser(
          description="GNN baselines on ogbgmol* data with Pytorch Geometrics"
      )
      parser.add_argument(
          "--device",
          type=int,
          default=0,
          help="which gpu to use if any (default: 0)",
      )
      parser.add_argument(
          "--num_mlp_layers",
          type=int,
          default=6,
          help="number of mlp layers (default: 6)",
      )
      parser.add_argument(
          "--drop_ratio",
          type=float,
          default=0.2,
          help="dropout ratio (default: 0.2)",
      )
      parser.add_argument(
          "--batch_size",
          type=int,
          default=256,
          help="input batch size for training (default: 256)",
      )
      parser.add_argument(
          "--emb_dim",
          type=int,
          default=1600,
          help="embedding dimensionality (default: 1600)",
      )
      parser.add_argument("--train_subset", action="store_true")
      parser.add_argument(
          "--epochs",
          type=int,
          default=100,
          help="number of epochs to train (default: 100)",
      )
      parser.add_argument(
          "--num_workers",
          type=int,
          default=0,
          help="number of workers (default: 0)",
      )
      parser.add_argument(
          "--radius", type=int, default=2, help="radius (default: 2)"
      )
      parser.add_argument(
          "--log_dir",
          type=str,
          default="",
          help="tensorboard log directory",
      )
      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")
      )
  
      dataset = PCQM4MDataset(root="dataset/", only_smiles=True)
      fp_processed_file = preprocess_fp(dataset, args.radius)
  
      data_dict = torch.load(fp_processed_file)
      X, Y = data_dict["X"], data_dict["Y"]
  
      split_idx = dataset.get_idx_split()
      ### automatic evaluator. takes dataset name as input
      evaluator = PCQM4MEvaluator()
  
      if args.train_subset:
          print("train subset")
          subset_ratio = 0.1
          subset_idx = torch.randperm(len(split_idx["train"]))[
              : int(subset_ratio * len(split_idx["train"]))
          ]
          train_dataset = TensorDataset(
              X[split_idx["train"][subset_idx]],
              Y[split_idx["train"][subset_idx]],
          )
  
      else:
          train_dataset = TensorDataset(
              X[split_idx["train"]], Y[split_idx["train"]]
          )
  
      valid_dataset = TensorDataset(
          X[split_idx["valid"]], Y[split_idx["valid"]]
      )
      test_dataset = TensorDataset(
          X[split_idx["test-dev"]], Y[split_idx["test"]]
      )
  
      train_loader = DataLoader(
          train_dataset,
          batch_size=args.batch_size,
          shuffle=True,
          num_workers=args.num_workers,
      )
      valid_loader = DataLoader(
          valid_dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.num_workers,
      )
  
      if args.save_test_dir != "":
          test_loader = DataLoader(
              test_dataset,
              batch_size=args.batch_size,
              shuffle=False,
              num_workers=args.num_workers,
          )
  
      if args.checkpoint_dir != "":
          os.makedirs(args.checkpoint_dir, exist_ok=True)
  
      model = MLP(
          num_mlp_layers=args.num_mlp_layers,
          emb_dim=args.emb_dim,
          drop_ratio=args.drop_ratio,
      ).to(device)
  
      num_params = sum(p.numel() for p in model.parameters())
      print(f"#Params: {num_params}")
  
      optimizer = optim.Adam(model.parameters(), lr=0.001)
  
      if args.log_dir != "":
          writer = SummaryWriter(log_dir=args.log_dir)
  
      best_valid_mae = 1000
  
      if args.train_subset:
          scheduler = StepLR(optimizer, step_size=300, gamma=0.25)
          args.epochs = 1000
      else:
          scheduler = StepLR(optimizer, step_size=30, gamma=0.25)
  
      for epoch in range(1, args.epochs + 1):
          print("=====Epoch {}".format(epoch))
          print("Training...")
          train_mae = train(model, device, train_loader, optimizer)
  
          print("Evaluating...")
          valid_mae = eval(model, device, valid_loader, evaluator)
  
          print({"Train": train_mae, "Validation": valid_mae})
  
          if args.log_dir != "":
              writer.add_scalar("valid/mae", valid_mae, epoch)
              writer.add_scalar("train/mae", train_mae, epoch)
  
          if valid_mae < best_valid_mae:
              best_valid_mae = valid_mae
              if args.checkpoint_dir != "":
                  print("Saving checkpoint...")
                  checkpoint = {
                      "epoch": epoch,
                      "model_state_dict": model.state_dict(),
                      "optimizer_state_dict": optimizer.state_dict(),
                      "scheduler_state_dict": scheduler.state_dict(),
                      "best_val_mae": best_valid_mae,
                      "num_params": num_params,
                  }
                  torch.save(
                      checkpoint,
                      osp.join(args.checkpoint_dir, "checkpoint.pt"),
                  )
  
              if args.save_test_dir != "":
                  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-dev",
                  )
  
          scheduler.step()
  
          print(f"Best validation MAE so far: {best_valid_mae}")
  
      if args.log_dir != "":
          writer.close()
  
imports: import torch
  from torch.utils.data import DataLoader, TensorDataset
  import torch.optim
  import torch.nn.functional
  from torch.utils.tensorboard import SummaryWriter
  from torch.optim.lr_scheduler import StepLR
  import os
  import os.path
  from tqdm.auto import tqdm
  import argparse
  import time
  import numpy
  import random
  from rdkit import Chem
  from rdkit.Chem import AllChem
  from ogb.lsc import PCQM4MDataset, PCQM4MEvaluator
label_desc: Use image node_id 4 for calling a global function with example usage: main_mlp() without return types