Fixed sample method of the model

anordertoreclaim · Aug 4, 2019 · 0224738 · 0224738
1 parent 03d5598
commit 0224738
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pixelcnn/model.py b/pixelcnn/model.py
@@ -142,7 +142,7 @@ def sample(self, shape, count, device='cuda'):
                     for c in range(channels):
                         unnormalized_probs = self.forward(samples)
                         pixel_probs = torch.softmax(unnormalized_probs[:, :, c, i, j], dim=1)
-                        sampled_levels = torch.multinomial(pixel_probs, 1).squeeze()
+                        sampled_levels = torch.multinomial(pixel_probs, 1).squeeze() / (self.color_levels - 1)
                         samples[:, c, i, j] = sampled_levels
 
         return samples