changed line of cuda to device

README.md

@@ -23,7 +23,8 @@ How to use the model?
 from diffusers import DiffusionPipeline
 
 wsi_generator = DiffusionPipeline.from_pretrained("kaveh/wsi_generator")
-wsi_generator.to("cuda")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+wsi_generator.to(device)
 
 generated_image = wsi_generator().images[0]
 generated_image.save("wsi_generated.png")