Fix pixelcnn test

juliuskunze · Nov 29, 2019 · ce6e2b6 · ce6e2b6
1 parent bd2104b
commit ce6e2b6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/examples/pixelcnn.py b/examples/pixelcnn.py
@@ -283,7 +283,7 @@ def main(batch_size=32, nr_filters=8, epochs=10, step_size=.001, decay_rate=.999
     opt = Adam(exponential_decay(step_size, 1, decay_rate))
     state = opt.init(loss.init_parameters(next(test_batches), key=init_key))
 
-    for epoch in range(0):
+    for epoch in range(epochs):
         for batch in get_train_batches():
             key, update_key = random.split(key)
             i = opt.get_step(state)

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -7,7 +7,7 @@
 from jax.random import PRNGKey
 
 from examples.mnist_vae import gaussian_sample, bernoulli_logpdf, gaussian_kl
-from examples.pixelcnn import PixelCNNPP
+from examples.pixelcnn import PixelCNNPP, image_dtype
 from examples.wavenet import calculate_receptive_field, discretized_mix_logistic_loss, Wavenet
 from jaxnet import parametrized, Dense, Sequential, Conv, flatten, GRUCell, Rnn, \
     Parameter, parameter, Reparametrized, L2Regularized, optimizers
@@ -244,7 +244,7 @@ def loss(batch):
 
 def test_pixelcnn():
     loss, _ = PixelCNNPP(nr_filters=1, nr_resnet=1)
-    images = np.zeros((2, 16, 16, 3), np.uint8)
+    images = np.zeros((2, 16, 16, 3), image_dtype)
     opt = optimizers.Adam()
     state = opt.init(loss.init_parameters(images, key=PRNGKey(0)))
     # take ~20s, disabled for faster tests: