Solved by b02202050 in post #2 
              
              
                Hi, 
one_hot_label = torch.cat((1 - current_label, current_label), dim=1)

explanation: 

(1 - current_label, current_label) would be (1,0) when current label is 0; and it would be (0,1) when current label is 1.
dim=1 means to operate on the 2nd dimension (which is channels)