Shape mismatch error during inference

[bc-bytes]

I used the following config to successfully train a binary segmentation model on my dataset (640x480 pixel images) using your code:

cfg = EasyDict()
cfg.batch_size = 2
cfg.epoch = 2000
cfg.learning_rate = 0.001
cfg.momentum = 0.9
cfg.weight_decay = 0.0001
cfg.patience = 10
cfg.inference_threshold = 0.75

cfg.transunet = EasyDict()
cfg.transunet.img_dim = 640
cfg.transunet.in_channels = 3
cfg.transunet.out_channels = 128
cfg.transunet.head_num = 4
cfg.transunet.mlp_dim = 640
cfg.transunet.block_num = 8
cfg.transunet.patch_dim = 16
cfg.transunet.class_num = 1

However, when I try to inference on the test set, I get this error:

einops.EinopsError: Error while processing rearrange-reduction pattern "b (x y ) c -> b c x y"
Input tensor shape: torch.Size([1, 1200, 1024]). Additional info: {'x': 40, 'y': 40}.
Shape mismatch, 1200 != 1600

Do you have any advice for this issue?