code: def main(config):
      config.models.gnn_saving_dir = os.path.join(
          os.path.dirname(os.path.abspath(__file__)), "checkpoints"
      )
      config.models.param = config.models.param[
          config.datasets.dataset_name
      ]
      print(OmegaConf.to_yaml(config))
      fix_random_seed(config.random_seed)
  
      if torch.cuda.is_available():
          device = torch.device("cuda", index=config.device_id)
      else:
          device = torch.device("cpu")
  
      dataset = get_dataset(
          dataset_root=config.datasets.dataset_root,
          dataset_name=config.datasets.dataset_name,
      )
      dataset.data.x = dataset.data.x.float()
      dataset.data.y = dataset.data.y.squeeze().long()
      if config.models.param.graph_classification:
          dataloader_params = {
              "batch_size": config.models.param.batch_size,
              "random_split_flag": config.datasets.random_split_flag,
              "data_split_ratio": config.datasets.data_split_ratio,
              "seed": config.datasets.seed,
          }
  
      model = get_gnnNets(
          dataset.num_node_features, dataset.num_classes, config.models
      )
      train_params = {
          "num_epochs": config.models.param.num_epochs,
          "num_early_stop": config.models.param.num_early_stop,
          "milestones": config.models.param.milestones,
          "gamma": config.models.param.gamma,
      }
      optimizer_params = {
          "lr": config.models.param.learning_rate,
          "weight_decay": config.models.param.weight_decay,
      }
  
      if config.models.param.graph_classification:
          trainer = TrainModel(
              model=model,
              dataset=dataset,
              device=device,
              graph_classification=config.models.param.graph_classification,
              save_dir=os.path.join(
                  config.models.gnn_saving_dir,
                  config.datasets.dataset_name,
              ),
              save_name=f"{config.models.gnn_name}_{len(config.models.param.gnn_latent_dim)}l",
              dataloader_params=dataloader_params,
          )
      else:
          trainer = TrainModel(
              model=model,
              dataset=dataset,
              device=device,
              graph_classification=config.models.param.graph_classification,
              save_dir=os.path.join(
                  config.models.gnn_saving_dir,
                  config.datasets.dataset_name,
              ),
              save_name=f"{config.models.gnn_name}_{len(config.models.param.gnn_latent_dim)}l",
          )
      trainer.train(
          train_params=train_params, optimizer_params=optimizer_params
      )
      _, _, _ = trainer.test()
  
imports: import os
  import hydra
  import torch
  import shutil
  import warnings
  from torch.optim import Adam
  from omegaconf import OmegaConf
  from benchmarks.xgraph.utils import check_dir
  from benchmarks.xgraph.gnnNets import get_gnnNets
  from benchmarks.xgraph.dataset import get_dataset, get_dataloader
  import torch.nn.functional
  from torch.optim.lr_scheduler import MultiStepLR
  from utils import fix_random_seed
label_desc: Use image node_id 1 for calling a global function with example usage: main(config) without return types