code: def training_step(self, batch, batch_idx):
      (x_q, x_k), _, _ = batch
  
      # update momentum
      update_momentum(self.backbone, self.backbone_momentum, 0.99)
      update_momentum(
          self.projection_head, self.projection_head_momentum, 0.99
      )
  
      # get queries
      q = self.backbone(x_q).flatten(start_dim=1)
      q = self.projection_head(q)
  
      # get keys
      k, shuffle = batch_shuffle(x_k)
      k = self.backbone_momentum(k).flatten(start_dim=1)
      k = self.projection_head_momentum(k)
      k = batch_unshuffle(k, shuffle)
  
      loss = self.criterion(q, k)
      self.log("train_loss_ssl", loss)
      return loss
  
imports: import copy
  import pytorch_lightning
  import torch
  import torch.nn
  import torchvision
  from lightly.data import LightlyDataset
  from lightly.loss import NTXentLoss
  from lightly.models import ResNetGenerator
  from lightly.models.modules.heads import MoCoProjectionHead
  from lightly.models.utils import batch_shuffle, batch_unshuffle, deactivate_requires_grad, update_momentum
  from lightly.transforms import MoCoV2Transform, utils
label_desc: Use image node_id 2 for calling the MocoModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss