fix condition block dtype mismatch in jit.save and enable 2 unitest

Author: HydrogenSulfate

deepmd/pd/utils/nlist.py

@@ -385,7 +385,11 @@ def build_multiple_neighbor_list(
         ).to(device=nlist.place)
         # nb x nloc x nsel
         nlist = paddle.concat([nlist, pad], axis=-1)
-        nsel = nsels[-1]
+        if paddle.is_tensor(nsel):
+            nsel = paddle.to_tensor(nsels[-1], dtype=nsel.dtype)
+        else:
+            nsel = nsels[-1]
+
     # nb x nall x 3
     coord1 = coord.reshape([nb, -1, 3])
     nall = coord1.shape[1]

source/tests/pd/model/test_jit.py

@@ -115,7 +115,6 @@ def tearDown(self):
         JITTest.tearDown(self)
 
 
-@unittest.skip("var dtype int32/int64 confused in if block")
 class TestEnergyModelDPA2(unittest.TestCase, JITTest):
     def setUp(self):
         input_json = str(Path(__file__).parent / "water/se_atten.json")

source/tests/pd/test_multitask.py

@@ -183,7 +183,6 @@ def tearDown(self):
                 shutil.rmtree(f)
 
 
-@unittest.skip("Paddle do not support MultiTaskSeA.")
 class TestMultiTaskSeA(unittest.TestCase, MultiTaskTrainTest):
     def setUp(self):
         multitask_se_e2_a = deepcopy(multitask_template)