code: def __init__(
      self,
      codebook_size,
      emb_dim,
      num_hiddens,
      straight_through=False,
      kl_weight=5e-4,
      temp_init=1.0,
  ):
      super().__init__()
      self.codebook_size = codebook_size  # number of embeddings
      self.emb_dim = emb_dim  # dimension of embedding
      self.straight_through = straight_through
      self.temperature = temp_init
      self.kl_weight = kl_weight
      self.proj = nn.Conv2d(
          num_hiddens, codebook_size, 1
      )  # projects last encoder layer to quantized logits
      self.embed = nn.Embedding(codebook_size, emb_dim)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from basicsr.utils import get_root_logger
  from basicsr.utils.registry import ARCH_REGISTRY
label_desc: Use image node_id 1 to create a new GumbelQuantizer object from inherited base classes: nn with example: obj = GumbelQuantizer(codebook_size, emb_dim, num_hiddens, straight_through, kl_weight, temp_init)