[PIR] Raise NotImplementedError in JIT save&load API (#61625)

SigureMo · web-flow · commit cc997ead69b1 · 2024-02-06T10:42:06.000+08:00
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -51,7 +51,7 @@
     dygraph_only,
 )
 from paddle.base.wrapped_decorator import wrap_decorator
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_mode, use_pir_api
 from paddle.nn import Layer
 from paddle.static.io import save_inference_model
 from paddle.utils.environments import (
@@ -914,6 +914,11 @@ def save(layer, path, input_spec=None, **configs):
             >>> save_function()
     """
 
+    if use_pir_api():
+        raise NotImplementedError(
+            "Currently, `paddle.jit.save` is not supported in PIR mode."
+        )
+
     # 1. input build & check
     prog_translator = ProgramTranslator()
     is_prim_infer = core._is_fwd_prim_enabled() and core._is_bwd_prim_enabled()
@@ -1508,6 +1513,10 @@ def load(path, **configs):
                 ...         print("Epoch {} batch {}: loss = {}".format(
                 ...             epoch_id, batch_id, np.mean(loss.numpy())))
     """
+    if use_pir_api():
+        raise NotImplementedError(
+            "Currently, `paddle.jit.load` is not supported in PIR mode."
+        )
     # 1. construct correct config
     config = _parse_load_config(configs)
     model_path, config = _build_load_path_and_config(path, config)
diff --git a/test/dygraph_to_static/test_fallback.py b/test/dygraph_to_static/test_fallback.py
@@ -117,11 +117,11 @@ def test_case_training(self):
         np.testing.assert_allclose(u_net(self.x).numpy(), [1, 1])
         assert u_net.training is False, "Training must be false."
 
-    @test_legacy_and_pt_and_pir
     def test_case_save_error(self):
         """
         test the save will raise error.
         """
+        # TODO(pir-save-load): Open this case after pir support save and load.
         u_net = UnsupportNet()
         u_net = paddle.jit.to_static(
             u_net, input_spec=[paddle.static.InputSpec(name='x', shape=[1])]
