tensorflow.python.framework.errors_impl.InternalError: Blas xGEMM launch failed

When running the image stylization code with tensorflow version 2.4.0, the following error occurred:

tensorflow.python.framework.errors_impl.InternalError: Blas xGEMM launch failed : a.shape=[1,480000,64], b.shape=[1,480000,64], m=64, n=64, k=480000 [Op:Einsum]

The following two solutions are found by consulting the data:
1. Add the following code to the program:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '/gpu:0'

The program can run normally, but the CPU is used, and the running speed of the program is much slower
2. Modify the cudnn version, but it is generally not recommended. It is too troublesome.

Read More: