In GPU-memory compression

I need to compress grayscale images in the GPU memory.

I tried encode_png((img * 255).type(torch.uint8).unsqueeze(dim=0)), however this does not support GPU.

I tried encode_jpeg((img * 255).type(torch.uint8).unsqueeze(dim=0)), however this does not seem to support grayscale images, despite what the documentation says:

  • input (Tensor [ channels , image_height , image_width ] or List [ Tensor [ channels , image_height , image_width ] ]) – (list of) uint8 image tensor(s) of c channels, where c must be 1 or 3

What would be the recommended way to do it?