Float vs. Int in torchvision.utils.save_image

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.

Internally the output array will be created via:

ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()

so you don’t need to multiply your predictions with 255 and can pass it as a float tensor.
The error is thrown at the .add_(0.5) step.

The documentation doesn’t mention any required ranges, so you might open a GitHub issue and link to this topic.

Thank you! Directly passing in mask_pred worked :slight_smile: Will open issue.

GitHub issue for reference: https://github.com/pytorch/vision/issues/1847