TypeError: Cannot subclass <class 'typing._SpecialForm'> while fine tuning GPT-J

Question

I am trying to fine tune GPT-J by following this GitHub Repository. When running the training command, I encounter this error:

Traceback (most recent call last):
  File "device_train.py", line 13, in <module>
    from mesh_transformer import util
  File "/home/shreyjain/mesh-transformer-jax/mesh_transformer/util.py", line 36, in <module>
    class ClipByGlobalNormState(OptState):
  File "/usr/lib/python3.8/typing.py", line 317, in __new__
    raise TypeError(f"Cannot subclass {cls!r}")
TypeError: Cannot subclass <class 'typing._SpecialForm'> 

This looks like a source code error but I am not sure. I have also raised an issue on GitHub regarding this. Any help will be appreciated!