Update app.py

app.py

@@ -7,6 +7,7 @@ from huggingface_hub import HfFileSystem
 from flax.serialization import msgpack_restore, from_state_dict
 import time
 from generator import Generator, LATENT_DIM
+import math
 
 generator = Generator()
 variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)