code: def main(args):
      # use CUDA_VISIBLE_DEVICES to select gpu
      device = (
          torch.device("cuda")
          if torch.cuda.is_available()
          else torch.device("cpu")
      )
      print("Using device:", device)
  
      dataset, train_data, val_data, test_data = get_dataset(
          args.dataset
      )
  
      g = dataset[0]
      g = g.to(device)
  
      input_size, representation_size = (
          g.ndata["feat"].size(1),
          args.graph_encoder_layer[-1],
      )
  
      # prepare transforms
      transform_1 = get_graph_drop_transform(
          drop_edge_p=args.drop_edge_p[0],
          feat_mask_p=args.feat_mask_p[0],
      )
      transform_2 = get_graph_drop_transform(
          drop_edge_p=args.drop_edge_p[1],
          feat_mask_p=args.feat_mask_p[1],
      )
  
      # scheduler
      lr_scheduler = CosineDecayScheduler(
          args.lr, args.lr_warmup_epochs, args.epochs
      )
      mm_scheduler = CosineDecayScheduler(1 - args.mm, 0, args.epochs)
  
      # build networks
      if args.dataset == "ppi":
          encoder = GraphSAGE_GCN(
              [input_size] + args.graph_encoder_layer
          )
      else:
          encoder = GCN([input_size] + args.graph_encoder_layer)
      predictor = MLP_Predictor(
          representation_size,
          representation_size,
          hidden_size=args.predictor_hidden_size,
      )
      model = BGRL(encoder, predictor).to(device)
  
      # optimizer
      optimizer = AdamW(
          model.trainable_parameters(),
          lr=args.lr,
          weight_decay=args.weight_decay,
      )
  
      # train
      for epoch in tqdm(
          range(1, args.epochs + 1), desc="  - (Training)  "
      ):
          train(
              epoch - 1,
              model,
              optimizer,
              lr_scheduler,
              mm_scheduler,
              transform_1,
              transform_2,
              g,
              args,
          )
          if epoch % args.eval_epochs == 0:
              val_scores, test_scores = eval(
                  model,
                  dataset,
                  device,
                  args,
                  train_data,
                  val_data,
                  test_data,
              )
              if args.dataset == "ppi":
                  print(
                      "Epoch: {:04d} | Best Val F1: {:.4f} | Test F1: {:.4f}".format(
                          epoch,
                          np.mean(val_scores),
                          np.mean(test_scores),
                      )
                  )
              else:
                  print(
                      "Epoch: {:04d} | Test Accuracy: {:.4f}".format(
                          epoch, np.mean(test_scores)
                      )
                  )
  
      # save encoder weights
      if not os.path.isdir(args.weights_dir):
          os.mkdir(args.weights_dir)
      torch.save(
          {"model": model.online_encoder.state_dict()},
          os.path.join(
              args.weights_dir, "bgrl-{}.pt".format(args.dataset)
          ),
      )
  
file_path: dgl/examples/pytorch/bgrl/main.py
imports: import copy
  import os
  import warnings
  import dgl
  import numpy
  import torch
  from eval_function import fit_logistic_regression, fit_logistic_regression_preset_splits, fit_ppi_linear
  from model import BGRL, compute_representations, GCN, GraphSAGE_GCN, MLP_Predictor
  from torch.nn.functional import cosine_similarity
  from torch.optim import AdamW
  from tqdm import tqdm
  from utils import CosineDecayScheduler, get_dataset, get_graph_drop_transform
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types