code: def train(
      step,
      model,
      optimizer,
      lr_scheduler,
      mm_scheduler,
      transform_1,
      transform_2,
      data,
      args,
  ):
      model.train()
  
      # update learning rate
      lr = lr_scheduler.get(step)
      for param_group in optimizer.param_groups:
          param_group["lr"] = lr
  
      # update momentum
      mm = 1 - mm_scheduler.get(step)
  
      # forward
      optimizer.zero_grad()
  
      x1, x2 = transform_1(data), transform_2(data)
  
      if args.dataset != "ppi":
          x1, x2 = dgl.add_self_loop(x1), dgl.add_self_loop(x2)
  
      q1, y2 = model(x1, x2)
      q2, y1 = model(x2, x1)
  
      loss = (
          2
          - cosine_similarity(q1, y2.detach(), dim=-1).mean()
          - cosine_similarity(q2, y1.detach(), dim=-1).mean()
      )
      loss.backward()
  
      # update online network
      optimizer.step()
      # update target network
      model.update_target_network(mm)
  
      return loss.item()
  
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 1 for calling a global function with example usage: train(step, model, optimizer, lr_scheduler, mm_scheduler, transform_1, transform_2, data, args) and returns: loss