Problems when trying to follow Readme to train an active inference.

[spring520]

I am using

1. Geforce 3060 Ti
2. tensorflow-gpu 2.4.0
3. python3.7
4. CUDA 11.0
5. cudnn 8.0

when trying to run train.py, the error messages show some problems with CUBLAS
failed to create cublasLt handle: CUBLAS_STATUS_INTERNAL_ERROR

I google this message, it looks like it's caused by the limitation of GPU memory. And I can't solve it with the solutions provided online. Can I ask what setting you use for the test? like which GPU and the version of tensorflow