Thanks @ptrblck for the quick reply.
If disable cudnn, the error doesn’t show up, but it’s running out of memory.
RuntimeError: CUDA out of memory. Tried to allocate 3.57 GiB (GPU 0; 23.62 GiB total capacity; 22.08 GiB already allocated; 691.56 MiB free; 5.43 MiB cached)

I am training a VNET with input image sizes [304, 304, 96]