Update causal_language_model_training_example.py

erfanzar · May 27, 2024 · 7e1bddb · 7e1bddb
1 parent 6ee12b8
commit 7e1bddb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/training/causal_language_model_training_example.py b/examples/training/causal_language_model_training_example.py
@@ -65,6 +65,7 @@ def main():
 
     dtype = jnp.bfloat16
     sharding_axis_dims = eval(FLAGS.sharding_axis_dims)
+    FLAGS.input_shape = eval(FLAGS.input_shape)
     qps = PartitionSpec(("dp", "fsdp"), "sp", None, "tp")
     kps = PartitionSpec(("dp", "fsdp"), "sp", None, "tp")
     vps = PartitionSpec(("dp", "fsdp"), "sp", None, "tp")