Update app.py

app.py

@@ -30,7 +30,7 @@ class CNN(nn.Module):
     return x
 
 cnn = CNN()
-params = cnn.init(jax.random.PRNGKey(0), jnp.ones([1, 50, 50, 3]))['params']
+params = cnn.init(jax.random.PRNGKey(0), jnp.ones([2, 50, 50, 3]))['params']
 
 fs = HfFileSystem(token=hf_key)
 with fs.open("PrakhAI/CatVsDog/checkpoint.msgpack", "rb") as f: