Update app.py

app.py

@@ -22,7 +22,7 @@ def infer(prompt, init_image):
     init_image = Image.open(init_image).convert("RGB")
     init_image = init_image.resize((128, 128))
 
-    with torch.autocast():
+    with torch.autocast('cuda'):
         res = pipe.train(
             prompt,
             init_image,