fix GPU test OOM problem

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/tf/env.py

@@ -483,6 +483,9 @@ def _get_package_constants(
 
 op_module = get_module("deepmd_op")
 op_grads_module = get_module("op_grads")
+# prevent OOM when using with other backends
+# tf.config doesn't work for unclear reason
+set_env_if_empty("TF_FORCE_GPU_ALLOW_GROWTH", "true", verbose=False)
 
 # FLOAT_PREC
 GLOBAL_TF_FLOAT_PRECISION = tf.dtypes.as_dtype(GLOBAL_NP_FLOAT_PRECISION)

source/tests/pt/conftest.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import pytest
+import torch
+
+
+@pytest.fixture(scope="package", autouse=True)
+def clear_cuda_memory(request):
+    yield
+    torch.cuda.empty_cache()