Modify to enable sampling

simplzyu · Jun 11, 2019 · 70018bc · 70018bc
1 parent 494ff85
commit 70018bc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/pixelsnail.py b/pixelsnail.py
@@ -383,6 +383,7 @@ def forward(self, input, condition=None):
             condition = F.one_hot(condition, self.n_class).permute(0, 3, 1, 2).float()
             condition = self.cond_resnet(condition)
             condition = F.interpolate(condition, scale_factor=2)
+            condition = condition[:, :, :height, :]
 
         for block in self.blocks:
             out = block(out, background, condition=condition)

diff --git a/vqvae.py b/vqvae.py
@@ -222,7 +222,9 @@ def decode(self, quant_t, quant_b):
 
     def decode_code(self, code_t, code_b):
         quant_t = self.quantize_t.embed_code(code_t)
+        quant_t = quant_t.permute(0, 3, 1, 2)
         quant_b = self.quantize_b.embed_code(code_b)
+        quant_b = quant_b.permute(0, 3, 1, 2)
 
         dec = self.decode(quant_t, quant_b)