mask_pred = net(imgs)

It is mapped to values of 0.96, 0.006, and so on. (because of normalization: x/ 255)

I’d like to see if this format was put in imshow() to see if it was well segmented

so I try…

imgs # size(1,1,640,640) = size(batch_size, channel, width, height)

```
mask_pred = net(imgs)[0]
mask_pred = (mask_pred >0.5).float()
np_img = mask_pred.cpu().detach().numpy() # GPU Tensor -> CPU Tensor -> numpy
np_img = np.transpose(np_img, (1, 2, 0)) # <class 'tuple'>: (640,640,1), dtype=float32
plt.imshow(np_img)
```

TypeError: Invalid dimensions for image data

how to solve it?