update test suits fiiting gpu platform

Author: refraction-ray

check_all.sh

@@ -8,6 +8,8 @@ echo "pylint check"
 pylint tensorcircuit tests examples/*.py
 echo "pytest check"
 pytest -n auto --cov=tensorcircuit -vv -W ignore::DeprecationWarning
+# for test on gpu machine, please set `export TF_FORCE_GPU_ALLOW_GROWTH=true` for tf
+# and `export XLA_PYTHON_CLIENT_PREALLOCATE=false` for jax to avoid OOM in testing
 echo "sphinx check"
 cd docs && sphinx-build source build/html && sphinx-build source -D language="zh" build/html_cn
 echo "all checks passed, congratulation! 💐"

tests/test_backends.py

@@ -300,7 +300,7 @@ def test_backend_methods_2(backend):
 def test_device_cpu_only(backend):
     a = tc.backend.ones([])
     dev_str = tc.backend.device(a)
-    assert dev_str == "cpu"
+    assert dev_str in ["cpu", "gpu:0"]
     tc.backend.device_move(a, dev_str)
 
 

tests/test_interfaces.py

@@ -297,6 +297,7 @@ def test_dlpack_transformation(backend):
             target_backend=b,
             enable_dlpack=True,
         )
+        ans = tc.interfaces.which_backend(ans).device_move(ans, "cpu")
         np.testing.assert_allclose(ans, np.ones([2]))