From 16d4d64d47707b9206615ca1c5ce8da9f61a9f57 Mon Sep 17 00:00:00 2001
From: Julio Perez <jperez@nvidia.com>
Date: Mon, 21 Nov 2022 17:47:28 -0500
Subject: [PATCH] add in environment variable to fix GPU memory allocation for
 tensorflow

---
 ci/test_unit.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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