kaveh commited on
Commit
3b99c79
1 Parent(s): 110e003

changed line of cuda to device

Browse files
Files changed (1) hide show
  1. README.md +2 -1
README.md CHANGED
@@ -23,7 +23,8 @@ How to use the model?
23
  from diffusers import DiffusionPipeline
24
 
25
  wsi_generator = DiffusionPipeline.from_pretrained("kaveh/wsi_generator")
26
- wsi_generator.to("cuda")
 
27
 
28
  generated_image = wsi_generator().images[0]
29
  generated_image.save("wsi_generated.png")
23
  from diffusers import DiffusionPipeline
24
 
25
  wsi_generator = DiffusionPipeline.from_pretrained("kaveh/wsi_generator")
26
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
27
+ wsi_generator.to(device)
28
 
29
  generated_image = wsi_generator().images[0]
30
  generated_image.save("wsi_generated.png")