The thing to notice is that the downsampling and upsampling blocks are fundamentally different, not just in what they do internally, but in their output: the downsampling blocks output two separate return values (one for the next downsampling step and one for the “skip” connection), but the upsampling blocks output only one tensor, right? So you need to index the output of the downsampling blocks to pick which of the two outputs you want. But it is a mistake to treat the output of the upsampling blocks in the same way: indexing it just ends up stripping off one of the dimensions of the single output tensor with the unhappy result that you are seeing.