In semantic segmentation ,How to transfer mask map to image

In semantic segmentation .I got the output, and the output size is (batch_size,22,256,256), I think this is a mask map about spital position. The question is How to transfer the mask map to a image.
The code link follow:
I think the code have some wrong. Please help me .
I write a neural net for semantic segmentation,and the loss function value descend to 0.05 .
The parameters of my net only 58M, but I do not know how to display the output image ,so for the result ,I have no idea .

Why do you think the code is wrong?
What kind of error did you encounter?

I think the question regarding visualizing the model’s output should be continued here.

Thank you for you help. I have referenced your answer,but the decode image is very worst.
May be my net have some questions