This “flattening” operation can be performed by:
x = torch.randn(2, 1, 32, 128, 128)
x = x.permute(0, 2, 1, 3, 4)
x = x.view(-1, *x.size()[2:])
print(x.shape)
> torch.Size([64, 1, 128, 128])

This could be an easy way to change the input format.
The alternative approach would be to open the volume, grab some slices, and return only these.
However, the logic inside the __getitem__ method would be a bit more complicated, as you would have to e.g. reuse the same volume to load the missing slices and would have to map the passed index somehow to this logic or use a custom sampler etc.