Inference with mps backend on M1 returns +/- Inf

I load a superresolution model with

device = torch.device('cpu')
model = torch.load("model_srresnet.pth",map_location=device)["model"]
image = Variable(torch.from_numpy(im_input / 255.0).float())
image =
with torch.inference_mode():
    output = model(image)

With device ‘cpu’ all works as expected the the inference output looks good, i.e. if I do output.min(), output.max() I see tensor(0.0403) tensor(1.0430). However, if I change ‘cpu’ to ‘mps’ the output min-max is tensor(inf, device=‘mps:0’) tensor(-inf, device=‘mps:0’).

Just more info: I am using torch version 1.13.1 on MacOS Ventura (13.0).

Problem solved, installed the latest nightly pytorch (torch 2.0.0.dev20230121) and it’s all good.