code: def forward(self, z):
      hard = self.straight_through if self.training else True
  
      logits = self.proj(z)
  
      soft_one_hot = F.gumbel_softmax(
          logits, tau=self.temperature, dim=1, hard=hard
      )
  
      z_q = torch.einsum(
          "b n h w, n d -> b d h w", soft_one_hot, self.embed.weight
      )
  
      # + kl divergence to the prior loss
      qy = F.softmax(logits, dim=1)
      diff = (
          self.kl_weight
          * torch.sum(
              qy * torch.log(qy * self.codebook_size + 1e-10), dim=1
          ).mean()
      )
      min_encoding_indices = soft_one_hot.argmax(dim=1)
  
      return z_q, diff, {"min_encoding_indices": min_encoding_indices}
  
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 2 for calling the GumbelQuantizer obj's underlying member method code with example usage: obj.forward(z) and returns: z_q,  diff,  dict