clean up

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Nov 9, 2024 · e4bac35 · e4bac35
1 parent 215efc6
commit e4bac35
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 33 deletions.
diff --git a/source/jax2tf_tests/test_nlist.py b/source/jax2tf_tests/test_nlist.py
@@ -1,30 +1,20 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import unittest
 
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
 
-from ...utils import (
-    DP_TEST_TF2_ONLY,
+from deepmd.jax.jax2tf.nlist import (
+    build_neighbor_list,
+    extend_coord_with_ghosts,
+)
+from deepmd.jax.jax2tf.region import (
+    inter2phys,
 )
 
-if DP_TEST_TF2_ONLY:
-    from deepmd.jax.jax2tf.nlist import (
-        build_neighbor_list,
-        extend_coord_with_ghosts,
-    )
-    from deepmd.jax.jax2tf.region import (
-        inter2phys,
-    )
-
-    dtype = tnp.float64
+dtype = tnp.float64
 
 
-@unittest.skipIf(
-    not DP_TEST_TF2_ONLY,
-    reason="TF2 conflicts with TF1",
-)
 class TestNeighList(tf.test.TestCase):
     def setUp(self):
         self.nf = 3

diff --git a/source/jax2tf_tests/test_region.py b/source/jax2tf_tests/test_region.py
@@ -1,29 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 
-import unittest
-
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
 
-from ...seed import (
-    GLOBAL_SEED,
-)
-from ...utils import (
-    DP_TEST_TF2_ONLY,
+from deepmd.jax.jax2tf.region import (
+    inter2phys,
+    to_face_distance,
 )
 
-if DP_TEST_TF2_ONLY:
-    from deepmd.jax.jax2tf.region import (
-        inter2phys,
-        to_face_distance,
-    )
+GLOBAL_SEED = 20241109
 
 
-@unittest.skipIf(
-    not DP_TEST_TF2_ONLY,
-    reason="TF2 conflicts with TF1",
-)
 class TestRegion(tf.test.TestCase):
     def setUp(self):
         self.cell = tnp.array(