Update app.py

app.py

@@ -5,9 +5,12 @@ from flax.serialization import msgpack_restore, from_state_dict
 fs = HfFileSystem()
 
 with fs.open("PrakhAI/HelloWorld/checkpoint.msgpack", "rb") as f:
-    state_dict = msgpack_restore(f.read())
-    print(state_dict)
-    print(state_dict.keys())
+    state = msgpack_restore(f.read())
+
+print(type(state))
+
+# logits = state.apply_fn({'params': state.params}, batch['image'])
+# print(logits)
 
 x = st.slider('Select a value')
 st.write(x, 'squared is', x * x)