densenet not training when using tf.contrib.layers.recompute_grad

[Sirius083]

I want to implement memory efficient densenet, following the code in
https://github.com/joeyearsley/efficient_densenet_tensorflow/blob/master/models/densenet_creator.py, the traing process is stuck at first epoch
I have just changed the add_layer part

        def add_layer(l):

            def _add_layer(l):
                shape = l.get_shape().as_list()
                in_channel = shape[3]
                with tf.variable_scope(name) as scope:
                    c = BatchNorm('bn1', l)
                    c = tf.nn.relu(c)
                    c = conv('conv1', c, self.growthRate, 1)
                    l = tf.concat([c, l], 3)
                return l
            
            if self.efficient:
                _add_layer = tf.contrib.layers.recompute_grad(_add_layer)
            
            return _add_layer(l)

also add the key word argument "efficient" to specify whether use the memory efficient version.
However the training process stucked.
Using tensorflow 1.9
tensorpack 0.9.1
Do I need to change other parts in the tensorpack?
Thanks in advance