This time the channels of the input to final and the layer definition do not match.
final takes 2048 in_channels, while the torch.cat operation concatenates t5x, where each one has 512 channels.
You could try to set the in_channels in final to 4096.