512x512 model fails to build

[mcicat]

When I try to initialize the model with image size set to 512.

ValueError: Exception encountered when calling layer 'decoder_block' (type DecoderBlock).
model = TransUNet(image_size=512, pretrain=True, num_classes=1)

The model fails during compilation with the following errors. I also tried to change the number of classes but it fails with the same erorr.

Traceback (most recent call last):
  File "/home/xxxx/work/MachineLearning/Training/KerasUnetCollection/Train.py", line 296, in <module>
    train_model(train_data, train_data_count, validation_data, validation_data_count, run_dir, learning_rate, img_target_width, img_target_height, batch_size, model_type)
  File "/home/xxxx/work/MachineLearning/Training/KerasUnetCollection/Train.py", line 196, in train_model
    model = TransUNet(image_size=shape[0],
  File "/home/xxxx/.local/lib/python3.10/site-packages/transunet/model.py", line 93, in TransUNet
    y = decoder_layers.DecoderCup(decoder_channels=decoder_channels, n_skip=n_skip)(y, features)
  File "/home/xxxx/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_file5u92zic2.py", line 42, in tf__call
    ag__.for_stmt(ag__.converted_call(ag__.ld(enumerate), (ag__.ld(self).blocks,), None, fscope), None, loop_body, get_state_1, set_state_1, ('x',), {'iterate_names': '(i, decoder_block)'})
  File "/tmp/__autograph_generated_file5u92zic2.py", line 38, in loop_body
    x = ag__.converted_call(ag__.ld(decoder_block), (ag__.ld(x),), dict(skip=ag__.ld(skip)), fscope)
  File "/tmp/__autograph_generated_fileo7vfvy5z.py", line 26, in tf__call
    ag__.if_stmt(ag__.ld(skip) is not None, if_body, else_body, get_state, set_state, ('x',), 1)
  File "/tmp/__autograph_generated_fileo7vfvy5z.py", line 21, in if_body
    x = ag__.converted_call(ag__.ld(tf).concat, ([ag__.ld(x), ag__.ld(skip)],), dict(axis=-1), fscope)
ValueError: Exception encountered when calling layer "decoder_cup" (type DecoderCup).

in user code:

    File "/home/xxxx/.local/lib/python3.10/site-packages/transunet/decoder_layers.py", line 95, in call  *
        x = decoder_block(x, skip=skip)
    File "/home/xxxx/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_fileo7vfvy5z.py", line 26, in tf__call
        ag__.if_stmt(ag__.ld(skip) is not None, if_body, else_body, get_state, set_state, ('x',), 1)
    File "/tmp/__autograph_generated_fileo7vfvy5z.py", line 21, in if_body
        x = ag__.converted_call(ag__.ld(tf).concat, ([ag__.ld(x), ag__.ld(skip)],), dict(axis=-1), fscope)

    ValueError: Exception encountered when calling layer 'decoder_block' (type DecoderBlock).

    in user code:

        File "/home/xxxx/.local/lib/python3.10/site-packages/transunet/decoder_layers.py", line 71, in call  *
            x = tf.concat([x, skip], axis=-1)

        ValueError: Dimension 1 in both shapes must be equal, but are 32 and 64. Shapes are [?,32,32] and [?,64,64]. for '{{node decoder_cup/decoder_block/concat}} = ConcatV2[N=2, T=DT_FLOAT, Tidx=DT_INT32](decoder_cup/decoder_block/up_sampling2d/resize/ResizeBilinear, Placeholder_1, decoder_cup/decoder_block/concat/axis)' with input shapes: [?,32,32,512], [?,64,64,512], [] and with computed input tensors: input[2] = <-1>.


    Call arguments received by layer 'decoder_block' (type DecoderBlock):
      • inputs=tf.Tensor(shape=(None, 16, 16, 512), dtype=float32)
      • skip=tf.Tensor(shape=(None, 64, 64, 512), dtype=float32)


Call arguments received by layer "decoder_cup" (type DecoderCup):
  • hidden_states=tf.Tensor(shape=(None, 16, 16, 768), dtype=float32)
  • features=['tf.Tensor(shape=(None, 64, 64, 512), dtype=float32)', 'tf.Tensor(shape=(None, 128, 128, 256), dtype=float32)', 'tf.Tensor(shape=(None, 256, 256, 64), dtype=float32)']