From 62007aeeca829462ad9cff5e179cf664c7914999 Mon Sep 17 00:00:00 2001 From: provos Date: Tue, 14 Feb 2023 16:50:38 -0800 Subject: [PATCH] Update vqvae.py correctly initialize the moving counter of codebook entry utilization. initializing this to 0 leads to divides by ~0 and numerically unstable embeddings at the beginning of training. --- vqvae.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vqvae.py b/vqvae.py index 099fa1915e..87f3286cc2 100755 --- a/vqvae.py +++ b/vqvae.py @@ -35,7 +35,7 @@ def __init__(self, dim, n_embed, decay=0.99, eps=1e-5): embed = torch.randn(dim, n_embed) self.register_buffer("embed", embed) - self.register_buffer("cluster_size", torch.zeros(n_embed)) + self.register_buffer("cluster_size", torch.ones(n_embed)) self.register_buffer("embed_avg", embed.clone()) def forward(self, input):