code: def training_step(self, batch, batch_idx) -> None:
      images, targets = batch[0], batch[1]
      features = self.model.forward(images).flatten(start_dim=1)
      if self.normalize:
          features = F.normalize(features, dim=1)
      features = features.to(self.feature_dtype)
      self._train_features.append(features.cpu())
      self._train_targets.append(targets.cpu())
  
file_path: lightly/lightly/utils/benchmarking/knn_classifier.py
imports: from typing import Tuple, Union
  import torch
  import torch.nn.functional
  from pytorch_lightning import LightningModule
  from torch import Tensor
  from torch.nn import Module
  from lightly.utils.benchmarking import knn_predict
  from lightly.utils.benchmarking.topk import mean_topk_accuracy
label_desc: Use image node_id 2 for calling the KNNClassifier obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) without return types