code: def main(model_cnf):
      yaml = YAML(typ="safe")
      model_cnf = yaml.load(Path(model_cnf))
      model_name, model_params, train_params, pretrain_params = (
          model_cnf["name"],
          model_cnf["model"],
          model_cnf["train"],
          model_cnf["pretrain"],
      )
      logger.info(f"Model name: {model_name}")
      logger.info(f"Model params: {model_params}")
      logger.info(f"Train params: {train_params}")
  
      if model_params["targets"] in [
          "mu",
          "homo",
          "lumo",
          "gap",
          "zpve",
      ]:
          model_params["output_init"] = nn.init.zeros_
      else:
          # 'GlorotOrthogonal' for alpha, R2, U0, U, H, G, and Cv
          model_params["output_init"] = GlorotOrthogonal
  
      logger.info("Loading Data Set")
      dataset = QM9(
          label_keys=model_params["targets"], edge_funcs=[edge_init]
      )
  
      # data split
      train_data, valid_data, test_data = split_dataset(
          dataset,
          num_train=train_params["num_train"],
          num_valid=train_params["num_valid"],
          shuffle=True,
          random_state=train_params["data_seed"],
      )
      logger.info(f"Size of Training Set: {len(train_data)}")
      logger.info(f"Size of Validation Set: {len(valid_data)}")
      logger.info(f"Size of Test Set: {len(test_data)}")
  
      # data loader
      train_loader = DataLoader(
          train_data,
          batch_size=train_params["batch_size"],
          shuffle=True,
          collate_fn=_collate_fn,
          num_workers=train_params["num_workers"],
      )
  
      valid_loader = DataLoader(
          valid_data,
          batch_size=train_params["batch_size"],
          shuffle=False,
          collate_fn=_collate_fn,
          num_workers=train_params["num_workers"],
      )
  
      test_loader = DataLoader(
          test_data,
          batch_size=train_params["batch_size"],
          shuffle=False,
          collate_fn=_collate_fn,
          num_workers=train_params["num_workers"],
      )
  
      # check cuda
      gpu = train_params["gpu"]
      device = (
          f"cuda:{gpu}"
          if gpu >= 0 and torch.cuda.is_available()
          else "cpu"
      )
  
      # model initialization
      logger.info("Loading Model")
      if model_name == "dimenet":
          model = DimeNet(
              emb_size=model_params["emb_size"],
              num_blocks=model_params["num_blocks"],
              num_bilinear=model_params["num_bilinear"],
              num_spherical=model_params["num_spherical"],
              num_radial=model_params["num_radial"],
              cutoff=model_params["cutoff"],
              envelope_exponent=model_params["envelope_exponent"],
              num_before_skip=model_params["num_before_skip"],
              num_after_skip=model_params["num_after_skip"],
              num_dense_output=model_params["num_dense_output"],
              num_targets=len(model_params["targets"]),
              output_init=model_params["output_init"],
          ).to(device)
      elif model_name == "dimenet++":
          model = DimeNetPP(
              emb_size=model_params["emb_size"],
              out_emb_size=model_params["out_emb_size"],
              int_emb_size=model_params["int_emb_size"],
              basis_emb_size=model_params["basis_emb_size"],
              num_blocks=model_params["num_blocks"],
              num_spherical=model_params["num_spherical"],
              num_radial=model_params["num_radial"],
              cutoff=model_params["cutoff"],
              envelope_exponent=model_params["envelope_exponent"],
              num_before_skip=model_params["num_before_skip"],
              num_after_skip=model_params["num_after_skip"],
              num_dense_output=model_params["num_dense_output"],
              num_targets=len(model_params["targets"]),
              extensive=model_params["extensive"],
              output_init=model_params["output_init"],
          ).to(device)
      else:
          raise ValueError(f"Invalid Model Name {model_name}")
  
      if pretrain_params["flag"]:
          torch_path = pretrain_params["path"]
          target = model_params["targets"][0]
          model.load_state_dict(torch.load(f"{torch_path}/{target}.pt"))
  
          logger.info("Testing with Pretrained model")
          predictions, labels = evaluate(device, model, test_loader)
          test_mae = mean_absolute_error(labels, predictions)
          logger.info(f"Test MAE {test_mae:.4f}")
  
          return
      # define loss function and optimization
      loss_fn = nn.L1Loss()
      opt = optim.Adam(
          model.parameters(),
          lr=train_params["lr"],
          weight_decay=train_params["weight_decay"],
          amsgrad=True,
      )
      scheduler = optim.lr_scheduler.StepLR(
          opt, train_params["step_size"], gamma=train_params["gamma"]
      )
  
      # model training
      best_mae = 1e9
      no_improvement = 0
  
      # EMA for valid and test
      logger.info("EMA Init")
      ema_model = copy.deepcopy(model)
      for p in ema_model.parameters():
          p.requires_grad_(False)
  
      best_model = copy.deepcopy(ema_model)
  
      logger.info("Training")
      for i in range(train_params["epochs"]):
          train_loss = train(device, model, opt, loss_fn, train_loader)
          ema(ema_model, model, train_params["ema_decay"])
          if i % train_params["interval"] == 0:
              predictions, labels = evaluate(
                  device, ema_model, valid_loader
              )
  
              valid_mae = mean_absolute_error(labels, predictions)
              logger.info(
                  f"Epoch {i} | Train Loss {train_loss:.4f} | Val MAE {valid_mae:.4f}"
              )
  
              if valid_mae > best_mae:
                  no_improvement += 1
                  if no_improvement == train_params["early_stopping"]:
                      logger.info("Early stop.")
                      break
              else:
                  no_improvement = 0
                  best_mae = valid_mae
                  best_model = copy.deepcopy(ema_model)
          else:
              logger.info(f"Epoch {i} | Train Loss {train_loss:.4f}")
  
          scheduler.step()
  
      logger.info("Testing")
      predictions, labels = evaluate(device, best_model, test_loader)
      test_mae = mean_absolute_error(labels, predictions)
      logger.info("Test MAE {:.4f}".format(test_mae))
  
imports: import copy
  from pathlib import Path
  import click
  import dgl
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.data.utils import Subset
  from logzero import logger
  from modules.dimenet import DimeNet
  from modules.dimenet_pp import DimeNetPP
  from modules.initializers import GlorotOrthogonal
  from qm9 import QM9
  from ruamel.yaml import YAML
  from sklearn.metrics import mean_absolute_error
  from torch.utils.data import DataLoader
label_desc: Use image node_id 7 for calling a global function with example usage: main(model_cnf) without return types