class ExampleLayer(paddle.nn.Layer):
    def __init__(self, hidd):
        super().__init__()
        self.fn = paddle.nn.Linear(hidd, hidd, bias_attr=False)
    def forward(self, x):
        for i in range(10):
            x = paddle.nn.functional.softmax(x,-1)
        x = x.cast("float32")
        x = self.func(x)
        return x
    def func(self, x):
        x = x + x
        return self.fn(x)

batch = 4096
hidd = 1024
dtype = "bfloat16"
x = paddle.rand([batch, hidd], dtype=dtype)
mylayer = ExampleLayer(hidd)
model = paddle.jit.to_static(
    mylayer,
    input_spec=[
        paddle.static.InputSpec(
            shape=[None, None], dtype=dtype),
    ])
# save to static model
save_path = "./checkpoints/infer"
paddle.jit.save(model, save_path)
print(f"static model has been to {save_path}")

from paddle.inference import Config
from paddle.inference import create_predictor
from paddle.inference import PrecisionType
model_dir = "checkpoints/"
model_file = model_dir + "/infer.pdmodel"
params_file = model_dir + "/infer.pdiparams"
config = Config(model_file, params_file)
config.enable_memory_optim()
gpu_precision = PrecisionType.Float32
config.enable_use_gpu(1000, 0, gpu_precision)
predictor = create_predictor(config)
result = predictor.run([x])

# 用基于装饰器的方式推理，帮助用户省略步骤2和步骤3的代码，达到同等效果。
mylayer = paddle.incubate.jit.inference(mylayer)
# 开启TensorRT后端
# mylayer = paddle.incubate.jit.inference(mylayer, with_trt=True)
# 进行推理，返回得到结果
decorator_result = mylayer(x)

        for block in self.blocks:
            x, y = block(x, y, c, mask)

   @paddle.incubate.jit.inference()
    def transformer_blocks(self, x,y,c,mask):
        for block in self.blocks:
            x, y = block(x, y, c, mask)
        return x, y