No, you have to make sure your data is on the CPU.
Also, even if you train your model on the GPU, all you have to do is shift the output tensor to the cpu and store the images. Here is a minimal Example
Once you have your tensor in CPU, another possibility is to apply Sigmoid to your output and estimate a threshold (the mid point for example) in order to save it as an binary image.
from torchvision.utils import save_image
img1 = torch.sigmoid(output) # output is the output tensor of your UNet, the sigmoid will center the range around 0.
# Binarize the image
threshold = (img1.min() + img1.max()) * 0.5
ima = torch.where(img1 > threshold, 0.9, 0.1)
Or you could try to “greyscale” the image…
img1 = torch.sigmoid(output)
min = img1.min()
max = img1.max()
img2 = 1./(max-min) * img1 + 1.*min / (min-max)