code: def forward(self, z: torch.Tensor, uses_ddp: bool = False):
      distances = (
          (z.reshape(-1, self.embedding_dim) ** 2).sum(
              dim=-1, keepdim=True
          )
          + (self.embeddings**2).sum(dim=-1)
          - 2 * z.reshape(-1, self.embedding_dim) @ self.embeddings.T
      )
  
      closest = distances.argmin(-1).unsqueeze(-1)
  
      quantized_indices = closest.reshape(
          z.shape[0], z.shape[1], z.shape[2]
      )
  
      one_hot_encoding = (
          F.one_hot(closest, num_classes=self.num_embeddings)
          .type(torch.float)
          .squeeze(1)
      )
  
      # quantization
      quantized = one_hot_encoding @ self.embeddings
      quantized = quantized.reshape_as(z)
  
      if self.training:
          n_i = torch.sum(one_hot_encoding, dim=0)
  
          if uses_ddp:
              dist.all_reduce(n_i)
  
          self.cluster_size = self.cluster_size * self.decay + n_i * (
              1 - self.decay
          )
  
          dw = one_hot_encoding.T @ z.reshape(-1, self.embedding_dim)
  
          if uses_ddp:
              dist.all_reduce(dw)
  
          ema_embed = self.ema_embed * self.decay + dw * (
              1 - self.decay
          )
  
          n = torch.sum(self.cluster_size)
  
          self.cluster_size = (
              (self.cluster_size + 1e-5)
              / (n + self.num_embeddings * 1e-5)
              * n
          )
  
          self.embeddings.data.copy_(
              ema_embed / self.cluster_size.unsqueeze(-1)
          )
          self.ema_embed.data.copy_(ema_embed)
  
      commitment_loss = F.mse_loss(
          quantized.detach().reshape(-1, self.embedding_dim),
          z.reshape(-1, self.embedding_dim),
          reduction="mean",
      )
  
      quantized = z + (quantized - z).detach()
  
      loss = commitment_loss * self.commitment_loss_factor
      quantized = quantized.permute(0, 3, 1, 2)
  
      output = ModelOutput(
          quantized_vector=quantized,
          quantized_indices=quantized_indices.unsqueeze(1),
          loss=loss,
      )
  
      return output
  
file_path: benchmark_VAE/src/pythae/models/vq_vae/vq_vae_utils.py
imports: import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from ..base.base_utils import ModelOutput
  from .vq_vae_config import VQVAEConfig
label_desc: Use image node_id 2 for calling the QuantizerEMA obj's underlying member method code with example usage: obj.forward(z, uses_ddp) and returns: output