Use mps if available.

[bakkot]

This works on my mps machine. (I.e., runs, generates the image, and does so with a speed indicative of using mps instead of the CPU.)

[LaurentMazare]

Great, I'm going to merge this then and will port the other examples to do the same, thanks for trying it out!