NGC Tensorflow | OP_REQUIRES failed at conv_ops.cc

Using the 20.09-tf2-py3 nvidia:tensorflow docker image I tried to use the pretrained NASNetMobile architecture from Keras pretrained models. It crashes with the following error:

2020-10-10 11:23:05.941484: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at conv_ops.cc:1115 : Not found: No algorithm worked!
Traceback (most recent call last):
  File "test.py", line 16, in <module>
    model.predict(batch)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 130, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 1599, in predict
    tmp_batch_outputs = predict_function(iterator)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 846, in _call
    return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)  # pylint: disable=protected-access
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 1848, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 1924, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 550, in call
    ctx=ctx)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.NotFoundError:  No algorithm worked!
         [[node NASNet/stem_conv1/Conv2D (defined at test.py:16) ]] [Op:__inference_predict_function_25681]

Function call stack:
predict_function

The minimum example on which it crashes:

import numpy as np
import tensorflow as tf

model = tf.keras.applications.NASNetMobile(input_shape=None, include_top=True, weights='imagenet', input_tensor=None, pooling=None, classes=1000)
batch = np.random.randint(low=0, high=256, size=(16,224,224,3))
batch = tf.keras.applications.nasnet.preprocess_input(batch)
model.predict(batch)

The above code works correctly on stock TF 2.3 but not on nvidia:tensorflow 20.9

2 Likes