update to_static code for darcy

Author: HydrogenSulfate

examples/darcy/darcy2d.py

@@ -13,12 +13,15 @@
 # limitations under the License.
 
 import numpy as np
+from paddle import fluid
 
 import ppsci
 from ppsci.utils import config
 from ppsci.utils import logger
 
 if __name__ == "__main__":
+    fluid.core._set_prim_all_enabled(True)
+
     args = config.parse_args()
     # set random seed for reproducibility
     ppsci.utils.misc.set_random_seed(42)

ppsci/arch/base.py

@@ -78,6 +78,9 @@ def split_to_dict(
         Returns:
             Dict[str, paddle.Tensor]: Dict contains tensor.
         """
+        # TODO: num_or_sections must > 1 in static, but 1 is allowed in dygraph.
+        if len(keys) == 1:
+            return {key: data_tensor for i, key in enumerate(keys)}
         data = paddle.split(data_tensor, len(keys), axis=axis)
         return {key: data[i] for i, key in enumerate(keys)}
 

ppsci/utils/expression.py

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable
-from typing import Union
-
+from paddle import jit
 from paddle import nn
 
 from ppsci.autodiff import clear
@@ -37,6 +35,7 @@ class ExpressionSolver(nn.Layer):
     def __init__(self):
         super().__init__()
 
+    @jit.to_static
     def forward(self, expr_dict, input_dict, model):
         output_dict = {k: v for k, v in input_dict.items()}