I predict a binary segmentation mask using an array with values 1 and 0, multiply this by 255, and try to save the resulting array as an image. However, I’m getting the following error:

Traceback (most recent call last):
File “test.py”, line 71, in 
torchvision.utils.save_image(predicted, path_ + idx[0])
File “C:\Users\CCL\Anaconda3\lib\site-packages\torchvision\utils.py”, line 107, in save_image
ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to(‘cpu’, torch.uint8).numpy()      RuntimeError: result type Float can’t be cast to the desired output type Int

Here is my code:
mask_pred = net(img)
predicted = torch.sigmoid(mask_pred) > 0.5
predicted = torch.tensor(predicted, dtype=torch.int32)*255
torchvision.utils.save_image(predicted, path_ + idx[0])

I’m using torch version 1.4.0 and torchvision version 0.5.0. The above code works on torch v1.1.0 and torchvision v0.3.0.