Your code seems to be written based on TensorFlow 1.X. Have you considered using TensorFlow 2.X or PyTorch?