It seems that Google has recently released the latest version of Tensorflow in an incomplete state that has not been completely backtested with the existing source. I have solved the problem by downgrading tensorflow from the current 2.17.0 to the previous stable 2.15.0 version as follows. I have also updated the github source.
# Delete the current tensorflow version
!pip uninstall --y tensorflow
# tensorflow version downgrade
!pip install tensorflow==2.15.0
# Check if downgrade was done
import tensorflow as tf
print(tf.__version__)
# Install package matching the downgraded version of TensorFlow
!apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2
# Install the compatible cuDNN version for CUDA 11.8
!apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2
# Uninstall existing TensorFlow and related packages
!pip uninstall -y -q tensorflow keras tensorflow-estimator tensorflow-text
# Install the required protobuf version
!pip install protobuf~=3.20.3
# Install TensorFlow Datasets
!pip install -q tensorflow_datasets
# Install TensorFlow 2.15.0 and TensorFlow Text 2.15.0
!pip install -q -U tensorflow-text==2.15.0
By doing this, you can complete the training task in Colab within 1 hour based on T4 GPU.