Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import jax
-from datasets import load_dataset
 import numpy as np
 import jax.numpy as jnp
 from flax.jax_utils import replicate