I am using tensorflow-gpu version 2,2 with cuda 10,2 and cudNN 8,0,5 . When I fit a Unet model from Segmentation models, with input shape (64,64,3) the fitting process throws an error:
`InternalError: cuDNN launch failure : input shape ([32,3,64,64])
[[node model_2/bn_data/FusedBatchNormV3 (defined at :12) ]] [Op:__inference_train_function_13058]
Function call stack:
train_function`
I have also tried memory growth
gpus = tensorflow.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
for gpu in gpus:
tensorflow.config.experimental.set_memory_growth(gpu, True)
logical_gpus = tensorflow.config.experimental.list_logical_devices(‘GPU’)
print(len(gpus), “Physical GPUs,”, len(logical_gpus), “Logical GPUs”)
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e)
May I get any suggestion to solve this problem?