code: def get_km_label(
      feat_dir: Path,
      km_dir: Path,
      label_dir: Path,
      split: str,
      num_rank: int,
      device: torch.device,
  ) -> None:
      r"""Predict the labels by the KMeans clustering model.
      Args:
          feat_dir (Path): The directory that stores the dumped features.
          km_dir (Path): The directory that stores the KMeans model.
          label_dir (Path): The directory to save the predicted labels.
          split (str): The split of data. Options: [``train``, ``valid``].
          num_rank (int): The number of ranks for multi-processing in feature extraction.
          device (torch.device): The location to allocate for PyTorch Tensors.
              Options: [``torch.device('cpu')``, torch.device('cuda')``].
      Returns:
          None
      """
      if not label_dir.exists():
          label_dir.mkdir()
  
      km_path = _get_model_path(km_dir)
      label_path = label_dir / f"label_{split}.pt"
      apply_kmeans = ApplyKmeans(km_path, device)
      with open(label_path, "w") as f:
          for rank in range(1, num_rank + 1):
              offset = 0
              feat_path, len_path = _get_feat_lens_paths(
                  feat_dir, split, rank, num_rank
              )
              feats = torch.load(feat_path)
              length = torch.load(len_path).int()
              assert feats.shape[0] == length.sum()
              labels = apply_kmeans(feats.to(device)).tolist()
              for i in range(length.shape[0]):
                  label = labels[offset : offset + length[i]]
                  offset += length[i]
                  f.write(" ".join(map(str, label)) + "\n")
      _LG.info("Finished predicting labels successfully")
  
file_path: audio/examples/hubert/utils/kmeans.py
imports: import logging
  from pathlib import Path
  from typing import Tuple
  import torch
  from sklearn.cluster import MiniBatchKMeans
  from torch import Tensor
  from .common_utils import _get_feat_lens_paths, _get_model_path
label_desc: Use image node_id 3 for calling a global function with example usage: get_km_label(feat_dir, km_dir, label_dir, split, num_rank, device) without return types