pytorch - 💡(How to fix) Fix torch.fmin with CUDA out=int16 casts inputs before computation, producing incorrect result

import torch

a = torch.tensor([100000], dtype=torch.int64)
b = torch.tensor([-1], dtype=torch.int64)

out_cpu = torch.empty(1, dtype=torch.int16)
out_gpu = torch.empty(1, dtype=torch.int16, device="cuda")

torch.fmin(a, b, out=out_cpu)
torch.fmin(a.cuda(), b.cuda(), out=out_gpu)

print("expected:", torch.fmin(a, b).to(torch.int16))
print("cpu:", out_cpu)
print("gpu:", out_gpu.cpu())

# extra checks
print("cuda without out:", torch.fmin(a.cuda(), b.cuda()).cpu())
print("cuda then cast:", torch.fmin(a.cuda(), b.cuda()).to(torch.int16).cpu())