Solved by ptrblck in post #4 
              
              
                You can just permute the dimensions: 
inputs = inputs.permute(0, 3, 2, 1)

or change your Dataset to return the input in the right shape.