Change default PE type for NeoX to `neox_rotary`

2022-03-19 00:26:04 -04:00 · 2022-03-19 00:26:04 -04:00 · c2c139e940
parent 85a4959efa
commit c2c139e940
1 changed files with 1 additions and 1 deletions
--- a/tpu_mtj_backend.py
+++ b/tpu_mtj_backend.py
@ -945,7 +945,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
            "n_vocab": 50432,
            "n_vocab_padding": 0,
            "norm": "doublelayernorm",
-            "pe": "rotary",
+            "pe": "neox_rotary",
            "pe_rotary_dims": 24,
            "seq": 2048,
            "cores_per_replica": 8,