KerasでU-Net実装中に起きた不可思議なエラー

前提・実現したいこと

KerasのFunctional APIでセマンティックセグメンテーションのU-Netを実装しようとしていた途中に起きたエラーです。
モデルを学習させ、テスト画像で推論しようとしていたときに以下のエラーメッセージが発生しました。
入力画像の大きさは512x512です。
実装は主にこのサイトをもとに行いました。

発生している問題・エラーメッセージ

WARNING:tensorflow:Model was constructed with shape (None, 512, 512, 3) for input Tensor("input_1:0", shape=(None, 512, 512, 3), dtype=float32), but it was called on an input with incompatible shape (32, 512, 3)

ValueError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:1150 predict_function  *
        outputs = self.distribute_strategy.run(
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:951 run  **
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2290 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2649 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:1125 predict_step  **
        return self(x, training=False)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py:927 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:719 call
        convert_kwargs_to_constants=base_layer_utils.call_context().saving)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:888 _run_internal_graph
        output_tensors = layer(computed_tensors, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py:886 __call__
        self.name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/input_spec.py:180 assert_input_compatibility
        str(x.shape.as_list()))

    ValueError: Input 0 of layer conv2d is incompatible with the layer: expected ndim=4, found ndim=3. Full shape received: [32, 512, 3]

該当のソースコード

入力画像を簡単なものに置き換えて試しましたが、同様のエラーが出ました。
簡単のため、そちらのコードを載せます。

python
1test = np.zeros((512, 512, 3), dtype=np.int32)
2predict = model.predict(test)

モデル自体のinput/outputの形状も正しいと思われます。

Model: "model"
Layer (type)                    Output Shape         Param #     Connected to                     
input_1 (InputLayer)            [(None, 512, 512, 3) 0                                            
(略)
conv2d_18 (Conv2D)              (None, 512, 512, 1)  65          conv2d_17[0][0]