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.weight**2).sum(dim=-1)
          - 2
          * z.reshape(-1, self.embedding_dim)
          @ self.embeddings.weight.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.weight
      quantized = quantized.reshape_as(z)
  
      commitment_loss = F.mse_loss(
          quantized.detach().reshape(-1, self.embedding_dim),
          z.reshape(-1, self.embedding_dim),
          reduction="mean",
      )
  
      embedding_loss = F.mse_loss(
          quantized.reshape(-1, self.embedding_dim),
          z.detach().reshape(-1, self.embedding_dim),
          reduction="mean",
      ).mean(dim=-1)
  
      quantized = z + (quantized - z).detach()
  
      loss = (
          commitment_loss * self.commitment_loss_factor
          + embedding_loss * self.quantization_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 Quantizer obj's underlying member method code with example usage: obj.forward(z, uses_ddp) and returns: output