Convert to tensor in predict.py

Former-commit-id: be1ad32304f2365c0c1db6b0e7fb835acc0fbfed
2019-12-04 12:43:04 +01:00 · 2019-12-04 12:43:04 +01:00 · 0d1cc25ae2
parent 5a7e934560
commit 0d1cc25ae2
1 changed files with 1 additions and 1 deletions
--- a/predict.py
+++ b/predict.py
@ -23,7 +23,7 @@ def predict_img(net,
    net.eval()

    ds = BasicDataset('', '', scale=scale_factor)
-    img = ds.preprocess(full_img)
+    img = torch.from_numpy(ds.preprocess(full_img))

    img = img.unsqueeze(0)
    img = img.to(device=device, dtype=torch.float32)