Solved by ptrblck in post #2 
              
              
                You should unsqueeze the statistics, as you are currently passing scalar values: 
normal = transforms.Normalize(mean.unsqueeze(0), std.unsqueeze(0))

This should avoid the error, which is caused by mean[:, None, None] inside Normalize.