TypeError: '<' not supported between instances of 'InputLayer' and 'Sequential'

I am trying to run model-optimization/tensorflow_model_optimization/g3doc/guide/pruning/pruning_with_keras.ipynb in python3.6.

File "train_pruning.py", line 90, in
callbacks=callbacks)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 727, in fit
use_multiprocessing=use_multiprocessing)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_arrays.py", line 675, in fit
steps_name='steps_per_epoch')
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_arrays.py", line 260, in model_iteration
callbacks._call_begin_hook(mode)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/keras/callbacks.py", line 250, in _call_begin_hook
self.on_train_begin()
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/keras/callbacks.py", line 366, in on_train_begin
callback.on_train_begin(logs)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py", line 53, in on_train_begin
self.prunable_layers = pruning_wrapper.collect_prunable_layers(self.model)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py", line 358, in collect_prunable_layers
layer for layer in model.submodules
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/module/module.py", line 193, in submodules
return tuple(self._flatten(predicate=_is_module))
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/module/module.py", line 339, in _flatten_module
for leaf_path, leaf in nest.flatten_with_tuple_paths(module_dict[key]):
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/util/nest.py", line 1327, in flatten_with_tuple_paths
flatten(structure, expand_composites=expand_composites)))
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/util/nest.py", line 263, in flatten
return _pywrap_tensorflow.Flatten(structure, expand_composites)
File "/home/jerry/anaconda3/envs/py3/lib/python3.6/site-packages/tensorflow_core/python/pywrap_tensorflow_internal.py", line 2651, in Flatten
return _pywrap_tensorflow_internal.Flatten(nested, expand_composites)
TypeError: '<' not supported between instances of 'InputLayer' and 'Sequential'

I think it's python3 can not recognize the type?
Does anyone know this problem?

Thanks,

Start by running model-optimization/tensorflow_model_optimization/g3doc/guide/pruning/pruning_with_keras.ipynb under Python 3.6 and compare the traceback through pruning_callbacks.py and pruning_wrapper.py. Determine why model.submodules fails while collecting prunable layers; done means the notebook completes training without this TypeError.