Update server.py

server.py

@@ -26,7 +26,7 @@ if not os.path.exists(result_path):
     result_path = download(url, root)
 
 
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = Dalle.from_pretrained(result_path)  # This will automatically download the pretrained model.
 model.to(device=device)
 model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)