Auto_scale_batch_size fails for to bigger batch sizes, cuDNN failure
🐛 Bug

I'm using a pre-trained ResNet50 on 224x224 images with 16-bit precision, I wanted to test the auto_scale_batch_size functionality.
The output in the terminal is the following:
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Using native 16bit precision.
Batch size 2 succeeded, trying batch size 4
Batch size 4 succeeded, trying batch size 8
Batch size 8 succeeded, trying batch size 16
Batch size 16 succeeded, trying batch size 32
Batch size 32 succeeded, trying batch size 64
Batch size 64 succeeded, trying batch size 128
Batch size 128 succeeded, trying batch size 256
All good till then. On batch size 256 the GPU's memory certainly is not sufficient and it fails on a 2d convolution.
The auto_scaling should be aborted at this point and the batch_size fixed to 128.
Instead the script fails with the message "RuntimeError: Unable to find a valid cuDNN algorithm to run convolution".
To Reproduce
I'm not doing anything exceptional:

Parsing arguments with "--auto_scale_batch_size", "true"
Initiating model, datamodule and trainer using the parsed arguments
trainer.tune(model, dm)

Question
Is this a problem on the Pytorch Lightning side not capturing the exception or is this anyhow linked to Cuda and installing a different version could be enough?
Environment
1x GeForce RTX 2080Ti
Ubuntu 20.10
Python 3.8.5
Pytorch Lightning 1.1.8
Pytorch 1.7.1
Cuda 11.2
Env created with miniconda, packages installed with pip