From 1b186b9f1f04b1151c31b8088b1eacb3142e7814 Mon Sep 17 00:00:00 2001
From: HJ <92386084+JJJYmmm@users.noreply.github.com>
Date: Wed, 27 Mar 2024 13:26:32 +0800
Subject: [PATCH] Fix variable naming errors in EMAVectorQuantizer

Fix variable naming errors, embedding_dim and n_embed
---
 taming/modules/vqvae/quantize.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/taming/modules/vqvae/quantize.py b/taming/modules/vqvae/quantize.py
index d75544e4..a0d14052 100644
--- a/taming/modules/vqvae/quantize.py
+++ b/taming/modules/vqvae/quantize.py
@@ -362,8 +362,8 @@ class EMAVectorQuantizer(nn.Module):
     def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
                 remap=None, unknown_index="random"):
         super().__init__()
-        self.codebook_dim = codebook_dim
-        self.num_tokens = num_tokens
+        self.codebook_dim = embedding_dim
+        self.num_tokens = n_embed
         self.beta = beta
         self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps)