Yes, TensorFlow 2.5.0 and later versions which are based on cuDNN 8.x does not have the 2 billion elements per tensor limitation that was present in cuDNN 7.x. Now, you can use TensorFlow to train and run models with tensors that contain more than 2 billion elements which will improve the memory efficiency as well as improve the performance.