update gpt_with_pir.py

PaddlePaddle · Nov 10, 2023 · e9fb452 · e9fb452
1 parent ff56f25
commit e9fb452
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 15 deletions.
diff --git a/test/auto_parallel/gpt_with_pir.py b/test/auto_parallel/gpt_with_pir.py
@@ -170,21 +170,10 @@ def test_dp_with_fused_linear(self):
         out_dp_ir = engine_dp_ir.fit(
             self.dataset, 3, batch_size=self.batch_size, log_freq=1
         )
-        # TODO(zhiqiu): fix accuracy problem and use array_equal to check it
-        np.testing.assert_allclose(
-            out_dp_prog.history["loss"][0],
-            out_dp_ir.history["loss"][0],
-            rtol=1e-5,
-            err_msg='pass {} has wrong results!, \nu={}\nv={}\ndiff={}'.format(
-                __class__,
-                out_dp_prog.history["loss"][0],
-                out_dp_ir.history["loss"][0],
-                out_dp_prog.history["loss"][0] - out_dp_ir.history["loss"][0],
-            ),
+
+        self.check_results(
+            out_dp_prog.history["loss"][0], out_dp_ir.history["loss"][0]
         )
-        # self.check_results(
-        #     out_dp_prog.history["loss"][0], out_dp_ir.history["loss"][0]
-        # )
 
     def test_mp(self):
         self.enable_pir(False)

diff --git a/test/cpp/pir/pattern_rewrite/drr_fuse_linear_test.cc b/test/cpp/pir/pattern_rewrite/drr_fuse_linear_test.cc
@@ -140,5 +140,5 @@ TEST(DrrTest, FusedLinear) {
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
-  EXPECT_EQ(program.block()->size(), 22u);
+  EXPECT_EQ(program.block()->size(), 23u);
 }