Program not working on GPU but works on CPU

Intermediate
16

Hmm…

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
   #torch_dtype=torch.float32 # too large
    torch_dtype=torch.float16
).to("cuda")

prompt = "A clear sunny landscape with mountains and a river"
image = pipe(prompt=prompt).images[0]
image.save("test_image.png")