Update app.py

app.py

@@ -30,11 +30,11 @@ class Generator(nn.Module):
     x = nn.tanh(x)
 
 generator = Generator()
-params = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)['params']
+variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
 
 fs = HfFileSystem()
 with fs.open("PrakhAI/DigitGAN/g_checkpoint.msgpack", "rb") as f:
-  g_state = from_state_dict(params, msgpack_restore(f.read()))
+  g_state = from_state_dict(variables, msgpack_restore(f.read()))
 
 def sample_latent(key):
   return jax.random.normal(key, shape=(1, LATENT_DIM))