Solved by ptrblck in post #4 
              
              
                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…