AttributeError: module 'tensorflow._api.v2.train' has no attribute 'Optimizer'

[alvynabranches]

Command

python IndicBERT/train/run_pretraining.py \
--input_file=/content/drive/MyDrive/6000Sentences/TrainingEnglish/input.txt \
--output_dir=/content/drive/MyDrive/6000Sentences/TrainingEnglish/mlm_data.txt \
--do_train=True \
--bert_config_file=/content/drive/MyDrive/6000Sentences/TrainingEnglish/output_tokenizer/tokenizer_config.json \
--train_batch_size=4096 \
--max_seq_length=5 \
--max_predictions_per_seq=4 \
--num_train_steps=1000000 \
--num_warmup_steps=10000 \
--learning_rate=5e-4 \
--save_checkpoints_steps=100000

Error

Traceback (most recent call last):
  File "/content/IndicBERT/train/run_pretraining.py", line 23, in <module>
    import optimization
  File "/content/IndicBERT/train/optimization.py", line 87, in <module>
    class AdamWeightDecayOptimizer(tf.train.Optimizer):
AttributeError: module 'tensorflow._api.v2.train' has no attribute 'Optimizer'