Transpose axes when making CustomDataset

Have you read this thread?