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

[Flax] Allow retraining from save checkpoint (#12559)

Browse Source 
* fix_torch_device_generate_test

* remove @

* finish
This commit is contained in:
Patrick von Platen
2021-07-07 14:43:43 +01:00
committed by GitHub
parent 1d6623c6a2
commit 7d321b7689
3 changed files with 22 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py
View File

@@ -427,7 +427,14 @@ if __name__ == "__main__":
rng = jax.random.PRNGKey(training_args.seed)
dropout_rngs = jax.random.split(rng, jax.local_device_count())
model = FlaxAutoModelForMaskedLM.from_config(config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype))
if model_args.model_name_or_path:
model = FlaxAutoModelForMaskedLM.from_pretrained(
model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
)
else:
model = FlaxAutoModelForMaskedLM.from_config(
config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
)
# Store some constant
num_epochs = int(training_args.num_train_epochs)