NotFoundError: Container TF-TRT does not exist.

[PolinaDemochkina]

Hello,
I'm trying to use the image_classification.py script to quantize a model, however, I get the following error: "tensorflow.python.framework.errors_impl.NotFoundError: Container TF-TRT does not exist. (Could not find resource: TF-TRT/TRTEngineOp_1) [Op:GetCalibrationDataOp]"

Using Tensorflow 2.0

Here's the code I use to create and save a test model:

import math
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import optimizers
from tensorflow.keras.datasets import cifar100
import os
import numpy as np

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

if tf.test.gpu_device_name():
    print('GPU found')
else:
    print("No GPU found")
    
(x_train, y_train), (x_test, y_test) = cifar100.load_data()
x_train = np.float32(x_train)[:64*(x_train.shape[0]//64)]
x_train /= 255.0
i_shape = x_train[0].shape

inputs = layers.Input(i_shape)
base_model = keras.models.Sequential([
    layers.Conv2D(128, 3, padding='same', strides=(2, 2)),
    layers.BatchNormalization(),
    layers.LeakyReLU(0.2),
    layers.Flatten(),
    layers.Dense(5*128),
    layers.Flatten(),
])(inputs)

prediction = layers.Dense(100, activation='softmax')(base_model)
model = keras.Model(inputs=inputs, outputs=prediction)
optimizer = optimizers.Adam(0.00001)
model.compile(optimizer, 'categorical_crossentropy', metrics=['accuracy'])

tf.saved_model.save(model, "/mnt/")

The command I use to run the script is:

python image_classification.py --use_synthetic --saved_model_dir /mnt/ --num_warmup_iterations 50  --num_calib_inputs 128 --display_every 10 --use_trt --precision INT8 --max_workspace_size $((2**32)) --batch_size 128  --mode benchmark

Here's the error I get:

Traceback (most recent call last):
  File "image_classification.py", line 413, in <module>
    optimize_offline=args.optimize_offline)
  File "image_classification.py", line 185, in get_graph_func
    input_fn, calib_files, num_calib_inputs//batch_size))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/compiler/tensorrt/trt_convert.py", line 992, in convert
    _save_calibration_table)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/compiler/tensorrt/trt_convert.py", line 921, in _for_each_trt_node
    fn(node)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/compiler/tensorrt/trt_convert.py", line 988, in _save_calibration_table
    _get_canonical_engine_name(node.name))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/compiler/tf2tensorrt/ops/gen_trt_ops.py", line 138, in get_calibration_data_op
    resource_name, name=name, ctx=_ctx)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/compiler/tf2tensorrt/ops/gen_trt_ops.py", line 187, in get_calibration_data_op_eager_fallback
    attrs=_attrs, ctx=_ctx, name=name)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.NotFoundError: Container TF-TRT does not exist. (Could not find resource: TF-TRT/TRTEngineOp_1) [Op:GetCalibrationDataOp]

[MachineJeff]

try
docker pull nvcr.io/nvidia/tensorflow:19.09-py3