pir of test_prim_simplenet_cinn (#67584)

xiaoguoguo626807 · web-flow · commit f8798eab8410 · 2024-08-21T10:24:43.000+08:00
diff --git a/test/prim/model/test_prim_simplenet_cinn.py b/test/prim/model/test_prim_simplenet_cinn.py
@@ -78,15 +78,29 @@ def train(self, use_prim):
     def check_prim(self, net, use_prim):
         if not use_prim:
             return
-        fwd_ops = [
-            op.type
-            for op in net.forward.get_concrete_program(self.data, self.dout)[1]
-            .train_program.block(0)
-            .ops
-        ]
-
-        # Ensure that batch_norm is splitted into small ops
-        self.assertTrue('batch_norm' not in fwd_ops)
+        if paddle.framework.use_pir_api():
+            fwd_ops = [
+                op.name()
+                for op in net.forward.get_concrete_program(
+                    self.data, self.dout
+                )[0]
+                .main_program.global_block()
+                .ops
+            ]
+
+            # Ensure that batch_norm is splitted into small ops
+            self.assertTrue('pd_op.batch_norm' not in fwd_ops)
+        else:
+            fwd_ops = [
+                op.type
+                for op in net.forward.get_concrete_program(
+                    self.data, self.dout
+                )[1]
+                .train_program.block(0)
+                .ops
+            ]
+            # Ensure that batch_norm is splitted into small ops
+            self.assertTrue('batch_norm' not in fwd_ops)
 
     def test_cinn_prim(self):
         if paddle.device.get_device() == "cpu":
