tensorflow: ValueError: Cannot iterate over a shape with unknown rank

1. System information

• OS Platform and Distribution: Ubuntu 20.04
• TensorFlow installation (pip package or built from source): pip (python 3.7)
• TensorFlow library (version, if pip package or github SHA, if built from source): tensorflow-2.4.1

2. Code

Provide code to help us reproduce your issues using one of the following options:

import sys
from pathlib import Path

import tensorflow as tf

# Specify the model.
saved_model_dir = Path('training/Model/admin/test2/1/exported-model/1/')

if saved_model_dir.exists():
    print(f'Converting model: {str(saved_model_dir)}')
else:
    print(f'Could not find model: {str(saved_model_dir)}')
    sys.exit(1)

# Convert the model.
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(str(saved_model_dir))
tflite_model = converter.convert()

# Save the model.
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

print('Ready.')

3. Failure after conversion

Error message:

Traceback (most recent call last):
  File "convert-to-tflite.py", line 17, in <module>
    tflite_model = converter.convert()
  File "/home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/lite.py", line 1947, in convert
    return super(TFLiteConverter, self).convert()
  File "/home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/lite.py", line 1304, in convert
    **converter_kwargs)
  File "/home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/convert.py", line 606, in toco_convert_impl
    input_tensors, output_tensors, *args, **kwargs)
  File "/home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/convert.py", line 497, in build_toco_convert_protos
    for dim in shape:
  File "/home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/python/framework/tensor_shape.py", line 861, in __iter__
    raise ValueError("Cannot iterate over a shape with unknown rank.")
ValueError: Cannot iterate over a shape with unknown rank.

4. Any other info / logs

Startup log:

2021-02-09 20:09:36.646605: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/Qt/5.14.1/lib
2021-02-09 20:09:36.646636: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
Converting model: training/Model/admin/test2/1/exported-model/1
2021-02-09 20:09:38.125156: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:38.125298: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/Qt/5.14.1/lib
2021-02-09 20:09:38.125310: W tensorflow/stream_executor/cuda/cuda_driver.cc:326] failed call to cuInit: UNKNOWN ERROR (303)
2021-02-09 20:09:38.125331: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (snowblower): /proc/driver/nvidia/version does not exist
2021-02-09 20:09:38.125557: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-02-09 20:09:38.125891: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
WARNING:tensorflow:From /home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/convert_saved_model.py:60: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
2021-02-09 20:09:38.792020: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:196] None of the MLIR optimization passes are enabled (registered 0 passes)
2021-02-09 20:09:38.858416: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 3593310000 Hz
2021-02-09 20:09:38.954802: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:39.889489: I tensorflow/core/grappler/devices.cc:69] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 0
2021-02-09 20:09:39.889750: I tensorflow/core/grappler/clusters/single_machine.cc:356] Starting new session
2021-02-09 20:09:39.890025: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:39.953928: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:928] Optimization results for grappler item: graph_to_optimize
  function_optimizer: function_optimizer did nothing. time = 0.004ms.
  function_optimizer: function_optimizer did nothing. time = 0ms.

WARNING:tensorflow:From /home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/lite/python/util.py:327: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.convert_variables_to_constants`
WARNING:tensorflow:From /home/thijs/.virtualenvs/tflite/lib/python3.7/site-packages/tensorflow/python/framework/convert_to_constants.py:856: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.extract_sub_graph`
2021-02-09 20:09:40.561332: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:41.338828: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:42.358175: I tensorflow/core/grappler/devices.cc:69] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 0
2021-02-09 20:09:42.358313: I tensorflow/core/grappler/clusters/single_machine.cc:356] Starting new session
2021-02-09 20:09:42.358532: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-09 20:09:42.426773: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:928] Optimization results for grappler item: graph_to_optimize
  function_optimizer: function_optimizer did nothing. time = 0.004ms.
  function_optimizer: function_optimizer did nothing. time = 0ms.