Update README.md
Browse files
README.md
CHANGED
|
@@ -72,6 +72,12 @@ def preprocess(image_path, ref_size=512):
|
|
| 72 |
# Run inference
|
| 73 |
inp, original_size = preprocess("input.jpg")
|
| 74 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 75 |
with torch.no_grad():
|
| 76 |
semantic, detail, matte = model(inp, True)
|
| 77 |
|
|
|
|
| 72 |
# Run inference
|
| 73 |
inp, original_size = preprocess("input.jpg")
|
| 74 |
|
| 75 |
+
import torch
|
| 76 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 77 |
+
|
| 78 |
+
model.to(device)
|
| 79 |
+
inp.to(device)
|
| 80 |
+
|
| 81 |
with torch.no_grad():
|
| 82 |
semantic, detail, matte = model(inp, True)
|
| 83 |
|