code: def extract_embeddings(
      model, data_loader, normalize_features: bool = False
  ):
      embeddings = []
      targets = []
      full_targets = []
  
      for x, y in tqdm(data_loader):
          x = x.cuda()
          B, C, H, W = x.shape
  
          # Forward input to model to capture last feature map
          y_hat = model.trunk.forward(x, ["lastMAP"])[0]
          if normalize_features:
              y_hat = F.normalize(y_hat, dim=-1)
  
          # Reshape the annotation to have same shape as feature map (B, H, W, C)
          b, h, w, c = y_hat.shape
          y_target = F.interpolate(
              y.float(), size=(h, w), mode="nearest"
          ).int()
          y_target = y_target.permute((0, 2, 3, 1))
          assert y_hat.shape[:-1] == y_target.shape[:-1]
          y = y.permute((0, 2, 3, 1))
          assert y.shape[:-1] == torch.Size([B, H, W])
  
          # Now, only keep the last dimension to build an table of
          # pairs (embedding, corresponding class)
          y_hat = y_hat.reshape(
              (b * h * w, -1)
          )  # TODO - do that after, only to create the index?
          y_target = y_target.reshape((b * h * w, -1))
          embeddings.append(y_hat.cpu())
          targets.append(y_target.cpu())
          full_targets.append(y.cpu())
  
      # Combine all the embedding and targets together
      embeddings = torch.cat(embeddings, dim=0)
      targets = torch.cat(targets, dim=0).squeeze(-1)
      full_targets = torch.cat(full_targets, dim=0).squeeze(-1)
      return embeddings, targets, full_targets
  
file_path: vissl/tools/knn_segmentation.py
imports: import sys
  from typing import Any, List
  import faiss
  import numpy
  import sklearn.metrics
  import torch
  import torch.nn.functional
  import torch.utils.data
  import torchvision.datasets
  import torchvision.transforms
  from tqdm import tqdm
  from vissl.data import dataset_catalog
  from vissl.models import build_model
  from vissl.utils.checkpoint import CheckpointLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.knn_utils import get_sorted_predictions
label_desc: Use image node_id 2 for calling a global function with example usage: extract_embeddings(model, data_loader, normalize_features) and returns: embeddings,  targets,  full_targets