Save int tensor [0,255] as image

Hi all,
I am trying to save output from learn.predict with save_image. But I am getting error.
My code

img_hr,a,b = learn.predict(fn)
img_hr.dtype
torch.int64

Error when saving the image

save_image(img_hr,'img.jpg')
RuntimeError                              Traceback (most recent call last)
Input In [186], in <cell line: 1>()
----> 1 save_image(img_hr,'img.jpg')

File ~/mambaforge/lib/python3.9/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
     24 @functools.wraps(func)
     25 def decorate_context(*args, **kwargs):
     26     with self.clone():
---> 27         return func(*args, **kwargs)

File ~/mambaforge/lib/python3.9/site-packages/torchvision/utils.py:154, in save_image(tensor, fp, format, **kwargs)
    152 grid = make_grid(tensor, **kwargs)
    153 # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
--> 154 ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
    155 im = Image.fromarray(ndarr)
    156 im.save(fp, format=format)

File ~/mambaforge/lib/python3.9/site-packages/fastai/torch_core.py:365, in TensorBase.__torch_function__(cls, func, types, args, kwargs)
    363 if cls.debug and func.__name__ not in ('__str__','__repr__'): print(func, types, args, kwargs)
    364 if _torch_handled(args, cls._opt, func): types = (torch.Tensor,)
--> 365 res = super().__torch_function__(func, types, args, ifnone(kwargs, {}))
    366 dict_objs = _find_args(args) if args else _find_args(list(kwargs.values()))
    367 if issubclass(type(res),TensorBase) and dict_objs: res.set_meta(dict_objs[0],as_copy=True)

File ~/mambaforge/lib/python3.9/site-packages/torch/_tensor.py:1142, in Tensor.__torch_function__(cls, func, types, args, kwargs)
   1139     return NotImplemented
   1141 with _C.DisableTorchFunction():
-> 1142     ret = func(*args, **kwargs)
   1143     if func in get_default_nowrap_functions():
   1144         return ret

RuntimeError: result type Float can't be cast to the desired output type Long

Any ideas on how to save the image?
@muellerzr tagging you because of excellent reply last time.

Hi,

This error arises from the fact that img_hr is in integer format, but save_image expects it to be a float tensor. Therefore, the data must be converted to float and re-scaled from the range [0, 255] to [0, 1], which can be done through dividing by 255.

save_image(img_hr/255, 'img.jpg')

Does that solve your problem?