PrakhAI commited on
Commit
c2df5e9
·
1 Parent(s): 8d2ee0d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -7,6 +7,10 @@ from flax.training import train_state # Useful dataclass to keep train state
7
  from flax import linen as nn # Linen API
8
  from huggingface_hub import HfFileSystem
9
  from flax.serialization import msgpack_restore, from_state_dict
 
 
 
 
10
 
11
 
12
 
@@ -31,7 +35,7 @@ cnn = CNN()
31
  params = cnn.init(jax.random.PRNGKey(0), jnp.ones([1, 28, 28, 1]))['params']
32
 
33
 
34
- fs = HfFileSystem()
35
  with fs.open("PrakhAI/HelloWorld/checkpoint.msgpack", "rb") as f:
36
  params = from_state_dict(params, msgpack_restore(f.read())["params"])
37
 
 
7
  from flax import linen as nn # Linen API
8
  from huggingface_hub import HfFileSystem
9
  from flax.serialization import msgpack_restore, from_state_dict
10
+ import os
11
+
12
+
13
+ hf_key = os.environ['HF_KEY']
14
 
15
 
16
 
 
35
  params = cnn.init(jax.random.PRNGKey(0), jnp.ones([1, 28, 28, 1]))['params']
36
 
37
 
38
+ fs = HfFileSystem(token=hf_key)
39
  with fs.open("PrakhAI/HelloWorld/checkpoint.msgpack", "rb") as f:
40
  params = from_state_dict(params, msgpack_restore(f.read())["params"])
41