I am new in python can anyone help me how can I save the prediction mask as image I want to save y_batch_pred as image

Based on your screenshot it seems you are using TensorFlow, not PyTorch, so I would recommend posting this question in the TensorFlow discussion board.

yeah I am using tensorflow ,Thanks for replying