I have two stage UNets. The first stage will output two stacks of image rounds:

an 8-round odd numbered images
an 8-round even numbered images.

I then stack them on top of each other: x = torch.cat((odd, even)).
Then I pass x into the second stage UNet: out = network(x).
Now my intuition is that I need to reshuffle “out” because they are still in order (8 odds, 8 evens), not (odd, even, odd, even, odd, even, odd, even).
But my colleague said the entire network and architecture will do the reshuffling, and I don’t need to do reshuffling myself.
Is this true? Why is that? He also mentioned the idea of permutation invariant/agnostic, but that kind of contradicts what he suggested here?