class DemoNet(nn.Layer):
    def __init__(self, mesh):
        super().__init__()
        self._mesh = mesh
        self.linear_0 = nn.Linear(IMAGE_SIZE, IMAGE_SIZE)
        self.linear_1 = nn.Linear(IMAGE_SIZE, CLASS_NUM)
        self.relu = nn.ReLU()
        # shard the weights of this layer
        self.linear_0.weight = dist.shard_tensor(
            self.linear_0.weight,
            self._mesh,
            [Shard(1)],
            stop_gradient=False,
        )
        self.linear_1.weight = dist.shard_tensor(
            self.linear_1.weight,
            self._mesh,
            [Shard(0)],
            stop_gradient=False,
        )
    def forward(self, x):
        out = self.linear_0(x)
        out = self.relu(out)
        out = self.linear_1(out)
        return out
mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
layer = DemoNet(mesh)
opt = paddle.optimizer.SGD(
    learning_rate=0.1, parameters=layer.parameters()
)
loss_fn = nn.MSELoss()
loader = create_data_loader()
dist_loader = dist.shard_dataloader(loader, meshes=[mesh])
dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
main_program = dist_model._engine._fwd_main_progs["train"]

{
    (%0) = "builtin.parameter" () {is_persistable:[true],parameter_name:"parameter_3",stop_gradient:[false]} : () -> builtin.tensor<8xf32>
    (%1) = "builtin.parameter" () {is_persistable:[true],parameter_name:"parameter_2",stop_gradient:[false]} : () -> builtin.tensor<16x8xf32>
    (%2) = "builtin.parameter" () {is_persistable:[true],parameter_name:"parameter_1",stop_gradient:[false]} : () -> builtin.tensor<16xf32>
    (%3) = "builtin.parameter" () {is_persistable:[true],parameter_name:"parameter_0",stop_gradient:[false]} : () -> builtin.tensor<16x16xf32>
    (%4) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"input0",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,16],stop_gradient:[true]} : () -> pd_dist.tensor<4x16xf32, mesh: {shape: [2], process_ids: [0,1], dim_names: [x]}, dims_mappings: [-1,-1]>
    (%5) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"label0",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,8],stop_gradient:[true]} : () -> pd_dist.tensor<4x8xf32, mesh: {shape: [2], process_ids: [0,1], dim_names: [x]}, dims_mappings: [-1,-1]>
    (%6) = "pd_op.matmul" (%4, %3) {stop_gradient:[false],transpose_x:false,transpose_y:false} : (pd_dist.tensor<4x16xf32, mesh: {shape: [2], process_ids: [0,1], dim_names: [x]}, dims_mappings: [-1,-1]>, builtin.tensor<16x16xf32>) -> builtin.tensor<4x16xf32>
    (%7) = "pd_op.add" (%6, %2) {stop_gradient:[false]} : (builtin.tensor<4x16xf32>, builtin.tensor<16xf32>) -> builtin.tensor<4x16xf32>
    (%8) = "pd_op.relu" (%7) {stop_gradient:[false]} : (builtin.tensor<4x16xf32>) -> builtin.tensor<4x16xf32>
    (%9) = "pd_op.matmul" (%8, %1) {stop_gradient:[false],transpose_x:false,transpose_y:false} : (builtin.tensor<4x16xf32>, builtin.tensor<16x8xf32>) -> builtin.tensor<4x8xf32>
    (%10) = "pd_op.add" (%9, %0) {stop_gradient:[false]} : (builtin.tensor<4x8xf32>, builtin.tensor<8xf32>) -> builtin.tensor<4x8xf32>
    (%11) = "pd_op.subtract" (%10, %5) {stop_gradient:[false]} : (builtin.tensor<4x8xf32>, pd_dist.tensor<4x8xf32, mesh: {shape: [2], process_ids: [0,1], dim_names: [x]}, dims_mappings: [-1,-1]>) -> builtin.tensor<4x8xf32>
    (%12) = "pd_op.square" (%11) {stop_gradient:[false]} : (builtin.tensor<4x8xf32>) -> builtin.tensor<4x8xf32>
    (%13) = "pd_op.mean" (%12) {axis:(pd_op.IntArray)[],keepdim:false,stop_gradient:[false]} : (builtin.tensor<4x8xf32>) -> builtin.tensor<f32>
}