Update README.md

README.md

@@ -40,6 +40,7 @@ pip install diffusers transformers accelerate scipy safetensors
 The image can then be generated with the following script.
 
 ```python
+import torch
 from diffusers import StableDiffusionPipeline
 
 keyword = "trpfrog"
@@ -49,7 +50,7 @@ model_id = "Prgckwb/trpfrog-diffusion"
 
 pipe = StableDiffusionPipeline.from_pretrained(
     model_id, torch_dtype=torch.float16
-).to(device)
+).to("cuda")
 
 image = pipe(prompt).images[0]
 image.save("trpfrog.jpg")