code: def __init__(
      self,
      *,
      model_name: str,
      encoder_projection_dropout: float,
      encoder_attention_dropout: float,
      encoder_ff_interm_dropout: float,
      encoder_dropout: float,
      encoder_layer_drop: float,
      mask_prob: float,
      mask_channel_prob: float,
      mask_channel_length: float,
      num_classes: int,
      aux_num_out: int,
      checkpoint: str,
      dataset_path: str,
      seconds_per_batch: float,
      subset: str,
      learning_rate: float,
      betas: Tuple[float, float],
      adam_eps: float,
      weight_decay: float,
      freeze_encoder_updates: int,
      warmup_updates: int,
      hold_updates: int,
      decay_updates: int,
  ):
      super().__init__()
  
      if model_name == "hubert_pretrain_base":
          self.model = torchaudio.models.hubert_pretrain_base(
              encoder_projection_dropout=encoder_projection_dropout,
              encoder_attention_dropout=encoder_attention_dropout,
              encoder_ff_interm_dropout=encoder_ff_interm_dropout,
              encoder_dropout=encoder_dropout,
              encoder_layer_drop=encoder_layer_drop,
              mask_prob=mask_prob,
              mask_channel_prob=mask_channel_prob,
              mask_channel_length=mask_channel_length,
              num_classes=num_classes,
          )
          self.aux = torch.nn.Linear(768, aux_num_out)
      elif model_name == "hubert_pretrain_large":
          self.model = torchaudio.models.hubert_pretrain_large(
              encoder_projection_dropout=encoder_projection_dropout,
              encoder_attention_dropout=encoder_attention_dropout,
              encoder_ff_interm_dropout=encoder_ff_interm_dropout,
              encoder_dropout=encoder_dropout,
              encoder_layer_drop=encoder_layer_drop,
              mask_prob=mask_prob,
              mask_channel_prob=mask_channel_prob,
              mask_channel_length=mask_channel_length,
              num_classes=num_classes,
          )
          self.aux = torch.nn.Linear(1024, aux_num_out)
      elif model_name == "hubert_pretrain_xlarge":
          self.model = torchaudio.models.hubert_pretrain_xlarge(
              encoder_projection_dropout=encoder_projection_dropout,
              encoder_attention_dropout=encoder_attention_dropout,
              encoder_ff_interm_dropout=encoder_ff_interm_dropout,
              encoder_dropout=encoder_dropout,
              encoder_layer_drop=encoder_layer_drop,
              mask_prob=mask_prob,
              mask_channel_prob=mask_channel_prob,
              mask_channel_length=mask_channel_length,
              num_classes=num_classes,
          )
          self.aux = torch.nn.Linear(1280, aux_num_out)
      else:
          raise ValueError(f"Unsupported model name: {model_name}.")
      self._load_checkpoint(checkpoint)
      for p in self.model.wav2vec2.feature_extractor.parameters():
          p.requires_grad = False
      self.loss_fn = torch.nn.CTCLoss(
          blank=0, reduction="sum", zero_infinity=True
      )
      self.optimizer = torch.optim.AdamW(
          list(self.aux.parameters()) + list(self.model.parameters()),
          lr=learning_rate,
          betas=betas,
          eps=adam_eps,
          weight_decay=weight_decay,
      )
      self.freeze_encoder_updates = freeze_encoder_updates
      self.lr_scheduler = TriStageLRScheduler(
          self.optimizer, warmup_updates, hold_updates, decay_updates
      )
      self.dataset_path = dataset_path
      self.seconds_per_batch = seconds_per_batch
      self.subset = subset
      self.automatic_optimization = False
      self.scaler = torch.cuda.amp.GradScaler()
  
imports: import math
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  import torchaudio
  import torchaudio.models.wav2vec2.components
  from dataset import _get_lengths_librilightlimited, _get_lengths_librispeech, BucketizeBatchSampler, CollateFnHubert, CollateFnLibriLightLimited, DistributedBatchSampler, HuBERTDataSet
  from lightning.pytorch import LightningModule
  from loss import hubert_loss
  from torch import Tensor
  from torch.optim.optimizer import Optimizer
  from torch.utils.data import DataLoader
label_desc: Use image node_id 1 to create a new HuBERTFineTuneModule object from inherited base classes: LightningModule with example: obj = HuBERTFineTuneModule()