SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

remove jax.ops.index (#16220)

Browse Source
This commit is contained in:
Suraj Patil
2022-03-17 17:51:43 +01:00
committed by GitHub
parent 8481ecefbd
commit 93d3fd8645
5 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/research_projects/jax-projects/model_parallel/README.md
View File

@@ -35,7 +35,7 @@ model = FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
emb = jnp.zeros((50264, model.config.hidden_size))
# update the first 50257 weights using pre-trained weights
emb = emb.at[jax.ops.index[:50257, :]].set(model.params["transformer"]["wte"]["embedding"])
emb = emb.at[:50257, :].set(model.params["transformer"]["wte"]["embedding"])
params = model.params
params["transformer"]["wte"]["embedding"] = emb