Change values inside a 3D tensor

I have a 224x224 binary image in a tensor (1, 224, 224), with 0-pixels representing background a 1-pixels representing foreground. I want to reshape it in a tensor (2, 224, 224), such as the first “layer” gt[0] has 1-pixels where there were 0-pixels in the original image and viceversa. This way one layer should show 1s where there is background and the other one will have 1s on the foreground (basically I need to have two complementary binary images in this tensor).

This is my code:

# gt is a tensor (1, 224, 224)
gt = gt.expand((2, 224, 224))  
backgr = gt[0]
foregr = gt[1]

backgr[backgr == 0] = 2 # swap all 0s in 1s and viceversa
backgr[backgr == 1] = 0
backgr[backgr == 2] = 1

gt[0] = backgr


The problem is both layers are modified with this code and I can’t figure out how to keep one of the two constant and change only gt[0].

You could use scatter_ to achieve this:

a = torch.zeros(1, 24, 24, dtype=torch.long).random_(2) # Your single channel map
b = torch.ones(1, 24, 24)
c = torch.zeros(2, 24, 24).scatter_(0, a, b)

Thank you! I actually found another solution, using:

gt = gt.repeat(2, 1, 1)

and then repeating the rest of the original code.