Cuda Version 9.0 with cudnn Version 7.6.5 error on Nvidia RTX A4000: failed to run cuBLAS routine cublasSgemm_v2: CUBLAS_STATUS_EXECUTION_FAILED

Accelerated Computing CUDA CUDA Setup and Installation
1

I running BERT BASE model on my Nvidia RTX A4000. I get this cuBLAS error when ever I run the model for fine-tuing.

failed to run cuBLAS routine cublasSgemm_v2: CUBLAS_STATUS_EXECUTION_FAILED

I am using:
GPU: Nvidia RTX A4000
Cuda: 9.0
Cudnn: 7.6.5
Tensorflow: tensorflow-gpu==1.10
python3.5

Anyone else with the same issue or a solution? Please help.