tensorflow.contrib error when using Python 3.8+

[Theelx]

OS: Ubuntu 20.04
Python: 3.8.8, 3.9.4
GPT-2-Simple: 0.7.2
Tensorflow: 2.2.2, 2.3.2, 2.4.1, 2.5.0rc1 (3.8.8); 2.5.0rc1 (3.9.4)

Code:

import gpt_2_simple as gpt2

Error (3.9.4):

Traceback (most recent call last):
  File "/root/test_gpt2.py", line 1, in <module>
    import gpt_2_simple as gpt2
  File "/root/.pyenv/versions/3.9.4/lib/python3.9/site-packages/gpt_2_simple/__init__.py", line 1, in <module>
    from .gpt_2 import *
  File "/root/.pyenv/versions/3.9.4/lib/python3.9/site-packages/gpt_2_simple/gpt_2.py", line 24, in <module>
    from gpt_2_simple.src import model, sample, encoder, memory_saving_gradients
  File "/root/.pyenv/versions/3.9.4/lib/python3.9/site-packages/gpt_2_simple/src/memory_saving_gradients.py", line 5, in <module>
    import tensorflow.contrib.graph_editor as ge
ModuleNotFoundError: No module named 'tensorflow.contrib'

I assume the issue is that I'm using tensorflow v2, which seems to lack this tensorflow.contrib module, however it's impossible to install any version of tensorflow v1 when using Python 3.8 or above. The page for this on PyPi says it supports any Python >= 3.6, and this broke on both Python 3.8 and 3.9, so one would assume there's a workaround for this, but so far I have been unable to find such a workaround.

Should the PyPi be updated to reflect that this won't work on Python 3.8+, should a workaround be included, or is there a workaround already that I have yet to find?

Edit: It seems tensorflow was purposely pinned to 1.15 in #200, so I assume PyPi needs to be told that this won't work on Python 3.8+.

[btarg]

I am also having this issue. How do I fix it?