Merge pull request #127 from boomb0om/master

fix incorrect argument pass
ai-forever · Jun 30, 2022 · 04dfbb1 · 04dfbb1
2 parents 3a7fe51 + 0a0e1bb
commit 04dfbb1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rudalle/vae/model.py b/rudalle/vae/model.py
@@ -99,7 +99,7 @@ def __init__(self, ddconfig, n_embed, embed_dim, dwt=False, kl_weight=1e-8):
     def encode(self, x, disable_gumbel_softmax=False):
         h = self.encoder(x)
         h = self.quant_conv(h)
-        quant, emb_loss, info = self.quantize(h, disable_gumbel_softmax)
+        quant, emb_loss, info = self.quantize(h, disable_gumbel_softmax=disable_gumbel_softmax)
         return quant, emb_loss, info
 
     def decode(self, quant):