The model could predict the classes as shown in the attached image; I was wondering how is it possible to display only the edges of the predicted objcet?

the image is the probability map of output.

```
outputs = model(t_image)
plt.figure()
prob = torch.exp(outputs)
prob_imgs = make_grid(prob.permute(1, 0, 2, 3))
prob_imgs = prob_imgs.detach()
plt.imshow(prob_imgs.permute(1, 2, 0))
plt.show()
```