torch.cdist does not work

[aced125]

Hi - torch cdist cuda is not implemented for fp16. How could I fix?

    return th.cdist(x, y, p=2)
  File "/home/ubuntu/anaconda3/envs/torch/lib/python3.7/site-packages/torch/functional.py", line 1119, in cdist
    return _VF.cdist(x1, x2, p, None)  # type: ignore
RuntimeError: "cdist_cuda" not implemented for 'Half'

[aced125]

This is not an issue in native pytorch AMP as it automatically knows to cast to fp32

[aced125]

As a workaround, I am doing:

import torch as th
import einops as ein


def cdist(x: th.Tensor, y: th.Tensor) -> th.Tensor:
    if x.dtype is th.float16 and x.is_cuda:
        x = ein.rearrange(x, "b l r -> b l () r")
        y = ein.rearrange(y, "b l r -> b () l r")
        return (x - y).norm(dim=-1, p=2)
    return th.cdist(x, y, p=2)

[Boltzmachine]

Maybe another option is to manually upcast it to float32 and cast it back?