CUBLAS_STATUS_INVALID_VALUE #3

Dear IgLM Team,

I'm excited to try out the model, but ran into the following error when running the README examples for iglm_infill, iglm_generate, and iglm_evaluate, as well as in the Python code example.

CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling cublasSgemmStridedBatched( handle, opa, opb, m, n, k, &alpha, a, lda, stridea, b, ldb, strideb, &beta, c, ldc, stridec, num_batches)

Environment information:
NVIDIA-SMI 470.94 Driver Version: 470.94 CUDA Version: 11.4
Cuda compilation tools, release 11.4, V11.4.152
Build cuda_11.4.r11.4/compiler.30521435_0
Python 3.8
cudatoolkit 11.3.1
nvidia-cudnn-cu11 8.5.0.96

Thanks for taking a look!

Best,
Chloe