错误：没有注册 OpKernel 以支持节点 DepthToSpace 使用的 Op 'DepthToSpace'

Question

I am trying to train a TensorFlow model. I get the following error:

Initializing models: 0%| | 0/5 [00:00<?, ?it/s]  
Error: No OpKernel was registered to support Op 'DepthToSpace' used by node DepthToSpace 
with these attrs: [data_format="NCHW", block_size=2, T=DT_FLOAT] 
Registered devices: [CPU] 
Registered kernels: 
 device='GPU'; T in [DT_QINT8] 
 device='GPU'; T in [DT_HALF] 
 device='GPU'; T in [DT_FLOAT] 
 device='CPU'; T in [DT_VARIANT]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_RESOURCE]; data_format in ["NHWC"]
 device='CPU'; T in [DT_STRING]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_BOOL]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_COMPLEX128]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_COMPLEX64]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_DOUBLE]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_FLOAT]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_BFLOAT16]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_HALF]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_INT32]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_INT8]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_UINT8]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_INT16]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_UINT16]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_UINT32]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_INT64]; data_format in ["NHWC"] 
 device='CPU'; T in [DT_UINT64]; data_format in ["NHWC"] [[DepthToSpace]] 

OS: Ubuntu 20.04, Python version: Python 3.8.5, Tensorflow version: 2.4.0, Graphics card: GeForce RTX 3090, NVIDIA driver: 460.67, CUDA version: 11.2
How do I fix it?

Answer 1

It is an issue related to TensorFlow GPU and CUDA. Just change TensorFlow version with Python version as following:

TensorFlow GPU version: 2.4.1
Python version: 3.6.9  

To test CUDA support for TensorFlow, just run following command:

import tensorflow as tf
tf.test.is_built_with_cuda()

It should return True