[fix]: Fix variable naming errors

2024-04-10 19:40:55 +08:00
parent 193b0398a8
commit e2492578c7
1 changed files with 2 additions and 2 deletions
@@ -365,8 +365,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)