Update app.py
Browse files
app.py
CHANGED
@@ -22,7 +22,7 @@ def infer(prompt, init_image):
|
|
22 |
init_image = Image.open(init_image).convert("RGB")
|
23 |
init_image = init_image.resize((128, 128))
|
24 |
|
25 |
-
with torch.autocast():
|
26 |
res = pipe.train(
|
27 |
prompt,
|
28 |
init_image,
|
|
|
22 |
init_image = Image.open(init_image).convert("RGB")
|
23 |
init_image = init_image.resize((128, 128))
|
24 |
|
25 |
+
with torch.autocast('cuda'):
|
26 |
res = pipe.train(
|
27 |
prompt,
|
28 |
init_image,
|