Shape error at II. Joint Optimization

[kafai7777]

I successfully run Shape Pre-Training but got error at II. Joint Optimization.

2023-02-23 00:27:17.819403: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:230] Shuffle buffer filled.
INFO:tensorflow:Error reported to Coordinator: in user code:

    /home/maojiahui/conda/nerfactor_b/nerfactor/nerfactor/models/nerfactor.py:209 call  *
        normal_pred = self._pred_normal_at(xyz)
    /home/maojiahui/conda/nerfactor_b/nerfactor/nerfactor/models/shape.py:203 chunk_func  *
        normals = out_layer(mlp_layers(surf_embed))
    /home/maojiahui/conda/nerfactor_b/nerfactor/nerfactor/models/shape.py:191 chunk_apply  *
        y_chunk = func(x_chunk)
    /home/maojiahui/conda/nerfactor_b/nerfactor/nerfactor/networks/mlp.py:46 __call__  *
        y = layer(x_)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py:1023 __call__  **
        self._maybe_build(inputs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py:2625 _maybe_build
        self.build(input_shapes)  # pylint:disable=not-callable
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/keras/layers/core.py:1198 build
        trainable=True)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py:655 add_weight
        caching_device=caching_device)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py:815 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:139 make_variable
        shape=variable_shape if variable_shape else None)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/distribute/shared_variable_creator.py:69 create_new_variable
        v = next_creator(**kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/distribute/distribute_lib.py:2111 creator_with_resource_vars
        created = self._create_variable(next_creator, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/distribute/mirrored_strategy.py:538 _create_variable
        distribute_utils.VARIABLE_POLICY_MAPPING, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/distribute/distribute_utils.py:306 create_mirrored_variable
        value_list = real_mirrored_creator(**kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/distribute/mirrored_strategy.py:530 _real_mirrored_creator
        v = next_creator(**kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py:752 variable_capturing_scope
        lifted_initializer_graph=lifted_initializer_graph, **kwds)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:264 __call__
        return super(VariableMetaclass, cls).__call__(*args, **kwargs)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py:293 __init__
        initial_value = initial_value()
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py:87 __call__
        self._checkpoint_position, shape, shard_info=shard_info)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py:122 __init__
        self.wrapped_value.set_shape(shape)
    /home/maojiahui/anaconda3/envs/nerfactor/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:1240 set_shape
        (self.shape, shape))

    ValueError: Tensor's shape (3, 128) is not compatible with supplied shape [63, 128]

I run this code with 3090Ti(tf2.5) . Looking forward to your help.

[kafai7777]

if statement was wrong.