code: def training_step(self, batch, batch_idx):
      if self.global_step > 0 and self.global_step % 300 == 0:
          # reset group features and weights every 300 iterations
          self._reset_group_features()
          self._reset_momentum_weights()
      else:
          # update momentum
          utils.update_momentum(
              self.backbone, self.backbone_momentum, 0.99
          )
          utils.update_momentum(
              self.projection_head, self.projection_head_momentum, 0.99
          )
  
      (x0, x1) = batch[0]
  
      if batch_idx % 2:
          # swap batches every second iteration
          x0, x1 = x1, x0
  
      x0_features = self.backbone(x0).flatten(start_dim=1)
      x0_encoded = self.projection_head(x0_features)
      x0_predicted = self.prediction_head(x0_encoded)
      x1_features = self.backbone_momentum(x1).flatten(start_dim=1)
      x1_encoded = self.projection_head_momentum(x1_features)
  
      # update group features and get group assignments
      assignments = self.smog.assign_groups(x1_encoded)
      group_features = self.smog.get_updated_group_features(x0_encoded)
      logits = self.smog(x0_predicted, group_features, temperature=0.1)
      self.smog.set_group_features(group_features)
  
      loss = self.criterion(logits, assignments)
  
      # use memory bank to periodically reset the group features with k-means
      self.memory_bank(x0_encoded, update=True)
  
      return loss
  
file_path: lightly/examples/pytorch_lightning/smog.py
imports: import copy
  import pytorch_lightning
  import torch
  import torchvision
  from sklearn.cluster import KMeans
  from torch import nn
  from lightly import loss, models
  from lightly.models import utils
  from lightly.models.modules import heads
  from lightly.transforms.smog_transform import SMoGTransform
label_desc: Use image node_id 5 for calling the SMoGModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss