Differing results in cpu/cuda compared to mps

mps_issues.py

import os

import numpy as np
import torch
from PIL import Image


def gather(consts: torch.Tensor, t: torch.Tensor):
    """Gather consts for $t$ and reshape to feature map shape"""
    c = consts.gather(-1, t)
    return c.reshape(-1, 1, 1, 1)


def img_to_tensor(im):
    return (
        torch.tensor(np.array(im.convert("RGB")) / 255, dtype=torch.float32)
        .permute(2, 0, 1)
        .unsqueeze(0)
        * 2
        - 1
    )


def tensor_to_image(t):
    return Image.fromarray(
        np.array(((t.squeeze().permute(1, 2, 0) + 1) / 2).clip(0, 1) * 255).astype(
            np.uint8
        )
    )


def q_xt_xtminus1(xtm1, t, beta):
    mean = gather(1.0 - beta, t) ** 0.5 * xtm1
    var = gather(beta, t)
    eps = torch.randn_like(xtm1)  # Noise shaped like xtm1
    return mean + (var**0.5) * eps


def add_noise(x, num_steps, device, save_every=1):
    x = x.to(device)
    beta = torch.linspace(0.0001, 0.04, num_steps).to(device)
    out_imgs = list()
    for t in range(num_steps):
        if t % save_every == 0:
            out_imgs.append(tensor_to_image(x.cpu()))

        t = torch.tensor(t, dtype=torch.long).to(device)
        x = q_xt_xtminus1(x, t, beta)

    return out_imgs


if __name__ == "__main__":
    num_steps = 100
    save_every = 20
    assert num_steps % save_every == 0, "num_steps must be divisible by save_every"
    # Download a sample from cifar10 dataset
    if not os.path.exists("cifar10.png"):
        os.system(
            "curl 'https://raw.githubusercontent.com/YoongiKim/CIFAR-10-images/master/train/horse/0003.jpg' > cifar10.png"
        )
    # Load image and convert to tensor of size 32x32
    im = Image.open("cifar10.png").resize((32, 32))
    x = img_to_tensor(im)
    mps_images = add_noise(x, num_steps, "mps", save_every)
    del x, im
    torch.cuda.empty_cache()
    ######### clear cache #########
    im = Image.open("cifar10.png").resize((32, 32))
    x = img_to_tensor(im)
    cpu_images = add_noise(x, num_steps, "cpu", save_every)
    # Save images
    image = Image.new("RGB", (32 * num_steps // save_every, 32 * 2))
    for i, (cpu_im, mps_im) in enumerate(zip(cpu_images, mps_images)):
        image.paste(cpu_im, (32 * i, 0))
        image.paste(mps_im, (32 * i, 32))

    image.save("cpu_mps.png")