add support for v3-32
Browse files- run_mlm_flax_stream.py +1 -0
run_mlm_flax_stream.py
CHANGED
|
@@ -591,6 +591,7 @@ if __name__ == "__main__":
|
|
| 591 |
)
|
| 592 |
|
| 593 |
# Setup train state
|
|
|
|
| 594 |
state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)
|
| 595 |
saved_step = -1
|
| 596 |
if model_args.model_name_or_path and "checkpoint" in model_args.model_name_or_path:
|
|
|
|
| 591 |
)
|
| 592 |
|
| 593 |
# Setup train state
|
| 594 |
+
print("TRAIN STATE")
|
| 595 |
state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)
|
| 596 |
saved_step = -1
|
| 597 |
if model_args.model_name_or_path and "checkpoint" in model_args.model_name_or_path:
|