Graph disconnected error when migrating to TensorFlow 2.1

Question

The code below is extracted from a DenseNet implementation in Keras. This code works fine on Keras 2.2.4 with TensorFlow 1.12.0 as the backend. When I run this same code on Tens