Colab notebook doesn't work

[shaneh2]

Hi,
I tried using the Colab notebook you linked to in your README for this repo at:
https://colab.research.google.com/drive/12u-CmB9evMIASNOqJtDW26gmNvSgepBv

However, the last code cell raises an exception:

2021-01-17 19:47:05.890096: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.10.1
Traceback (most recent call last):
  File "language_generation.py", line 16, in <module>
    import model_utils
  File "/content/XLNet-gen/model_utils.py", line 292, in <module>
    class AdamWeightDecayOptimizer(tf.train.Optimizer):
AttributeError: module 'tensorflow._api.v2.train' has no attribute 'Optimizer'