code: def main():
      args = get_args()
      set_seed_(args.seed)
  
      dense_model_name = args.dense_model_name
      quantizer_model_name = "kmeans"
      vocab_size = args.vocab_size
  
      # NB: Hubert is not serializable as-is, so to have a multi-worker dataloader
      # we have a worker-around: load the actual checkpoint on the first call - which
      # will happen in a worker process already. This behavior is enabled with
      # the `lazy_load` flag.
      dense_model = dispatch_dense_model(
          dense_model_name, lazy_load=True
      )
      quantizer_model = dispatch_quantizer(
          dense_model_name, quantizer_model_name, vocab_size
      )
  
      speech_encoder = SpeechEncoder(
          dense_model,
          quantizer_model,
          deduplicate=False,
          need_f0=False,
          add_bos_eos=True,
      )
  
      dataset = QuantizedLibriSpeech(
          speech_encoder,
          root="datasets",
          url="dev-clean",
          download=True,
          device="auto"
          # when we set `device` to auto, the dataset instance will check if it is
          # running within a worker process of a dataloader. If it is the case,
          # it will move SpeechEncoder to one of the available GPUs, depending on the
          # worker id. This way we can pack quite a few (GPU-hungry) Hubert instances running across
          # all GPUs in parallel, within the same standard DataLoader.
      )
  
      speaker_mapping = SpeakerDatasetWrapper.get_speaker_ids(
          dataset.dataset._walker
      )
      max_speaker_id = max(speaker_mapping.values())
      dataset = SpeakerDatasetWrapper(dataset, speaker_mapping)
  
      valid_size = int(0.1 * len(dataset))
      train_size = len(dataset) - valid_size
      train_data, valid_data = torch.utils.data.random_split(
          dataset,
          [train_size, valid_size],
          generator=torch.Generator().manual_seed(42),
      )
  
      train_loader = DataLoader(
          train_data,
          batch_size=args.batch_size,
          shuffle=True,
          collate_fn=dataset.collater,
          num_workers=4,
      )
      valid_loader = DataLoader(
          valid_data,
          batch_size=args.batch_size,
          shuffle=False,
          collate_fn=dataset.collater,
          num_workers=4,
      )
  
      if args.model_type == "baseline":
          model = ConstantBaseline(total_speakers=max_speaker_id + 1)
      elif args.model_type == "discrete":
          model = DiscreteClassifier(
              vocab_size=args.vocab_size
              + 3,  # accounting for bos, pad, eos
              embedding_size=32,
              n_heads=4,
              hidden_size=128,
              n_layers=2,
              dropout=0.1,
              pad_value=dataset.quantized_data.unit_pad,
              total_speakers=max_speaker_id + 1,
          )
      elif args.model_type == "continuous":
          input_size = {
              "hubert-base-ls960": 768,
              "cpc-big-ll6k": 512,
          }[dense_model_name]
  
          model = ContinuousClassifier(
              embedding_size=32,
              input_size=input_size,
              n_heads=4,
              hidden_size=128,
              n_layers=2,
              dropout=0.1,
              pad_value=dataset.quantized_data.unit_pad,
              total_speakers=max_speaker_id + 1,
          )
      else:
          assert False, "unknown model type"
  
      model.cuda()
      train(model, train_loader, valid_loader, args)
  
file_path: textlesslib/examples/speaker_probing/train.py
imports: import torch
  import argparse
  from textless.data.quantized_datasets import QuantizedLibriSpeech
  from torch.utils.data import DataLoader
  import torch.nn.functional
  from probes import ContinuousClassifier, DiscreteClassifier, ConstantBaseline
  from textless.data.speech_encoder import SpeechEncoder
  from textless import dispatch_dense_model, dispatch_quantizer
label_desc: Use image node_id 7 for calling a global function with example usage: main() without return types