I have a tensor of size (24, 2, 224, 224) in Pytorch.
24 = batch size
2 = matrixes representing foreground and background
224 = image height dimension
224 = image width dimension
This is the output of a CNN that performs binary segmentation. In each cell of the 2 matrixes is stored the probability for that pixel to be foreground or background:
[n][0][h][w] + [n][1][h][w] = 1 for every coordinate
I want to reshape it into a tensor of size (24, 1, 224, 224). The values in the new layer sholud be 0 or 1 according to the matrix in which the probability was higher.
How can I do that? Which function should I use?