@albanD Indeed! But in the broken code above, the first snippet, I also concatenate on dim=1.
ETA: to make it more visible
I do torch.cat([y_sample, z_sample]), dim=1, and just before that   print(y_sample.shape, z_sample.shape) outputs torch.Size([100, 10]) torch.Size([100, 32]).