Update app.py

app.py

@@ -28,6 +28,7 @@ class Generator(nn.Module):
     x = nn.relu(x)
     x = nn.ConvTranspose(features=1, kernel_size=(2, 2), strides=(2, 2))(x)
     x = nn.tanh(x)
+    return x
 
 generator = Generator()
 variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
@@ -42,8 +43,5 @@ def sample_latent(key):
 if st.button('Generate Digit'):
   latents = sample_latent(jax.random.PRNGKey(int(1_000_000 * time.time())))
   g_out = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
-  st.write(g_state['params'])
-  st.write(g_state['batch_stats'])
-  st.write(g_out)
   img = ((np.array(g_out)+1)*255./2.).astype(np.uint8)[0]
   st.image(Image.fromarray(np.repeat(img, repeats=3, axis=2)))