[PIR] fix xpu case (#67716)

wanghuancoder · web-flow · commit d61270198df7 · 2024-08-27T17:10:12.000+08:00
* fix xpu case
diff --git a/test/xpu/collective_reduce_api.py b/test/xpu/collective_reduce_api.py
@@ -59,7 +59,8 @@ def get_model(self, main_prog, startup_program, rank, dtype='float32'):
             tindata = paddle.static.data(
                 name="tindata", shape=[-1, 10, 1000], dtype=dtype
             )
-            tindata.desc.set_need_check_feed(False)
+            if not paddle.framework.use_pir_api():
+                tindata.desc.set_need_check_feed(False)
             paddle.distributed.reduce(tindata, dst=0)
             return [tindata]
 
@@ -75,7 +76,8 @@ def get_model_new(
             tindata = paddle.static.data(
                 name="tindata", shape=[10, 1000], dtype=dtype
             )
-            tindata.desc.set_need_check_feed(False)
+            if not paddle.framework.use_pir_api():
+                tindata.desc.set_need_check_feed(False)
             reduce_new(tindata, dst=0, reduce_type=reduce_type)
             return [tindata]
 
diff --git a/test/xpu/test_collective_reduce_xpu.py b/test/xpu/test_collective_reduce_xpu.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
+os.environ['FLAGS_enable_pir_api'] = '0'
+
 import unittest
 
 from get_test_cover_info import get_xpu_op_support_types
