diff --git a/ci/test_unit.sh b/ci/test_unit.sh
index 071ad5ca8..82c2a721d 100755
--- a/ci/test_unit.sh
+++ b/ci/test_unit.sh
@@ -18,6 +18,6 @@
 set -e
 
 # Run tests
-pytest -rsx tests/unit tests/notebook
+TF_GPU_ALLOCATOR=cuda_malloc_async pytest -rsx tests/unit tests/notebook