code: def training_step(self, batch, batch_idx):
      momentum = cosine_schedule(self.current_epoch, 10, 0.996, 1)
      update_momentum(self.backbone, self.backbone_momentum, m=momentum)
      update_momentum(
          self.projection_head,
          self.projection_head_momentum,
          m=momentum,
      )
      (x0, x1) = batch[0]
      p0 = self.forward(x0)
      z0 = self.forward_momentum(x0)
      p1 = self.forward(x1)
      z1 = self.forward_momentum(x1)
      loss = 0.5 * (self.criterion(p0, z1) + self.criterion(p1, z0))
      return loss
  
file_path: lightly/examples/pytorch_lightning/byol.py
imports: import copy
  import pytorch_lightning
  import torch
  import torchvision
  from torch import nn
  from lightly.loss import NegativeCosineSimilarity
  from lightly.models.modules import BYOLPredictionHead, BYOLProjectionHead
  from lightly.models.utils import deactivate_requires_grad, update_momentum
  from lightly.transforms.byol_transform import BYOLTransform, BYOLView1Transform, BYOLView2Transform
  from lightly.utils.scheduler import cosine_schedule
label_desc: Use image node_id 4 for calling the BYOL obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss