Unset train_global_input_shapes if always_use_train_for_model_init is not set, instead of making an assertion, for backward compatibility.
Unset train_global_input_shapes if always_use_train_for_model_init is not set,
instead of making an assertion, for backward compatibility.

PiperOrigin-RevId: 486693970

google/paxml,google/paxml