Update example_inference.py
Browse files- example_inference.py +1 -4
example_inference.py
CHANGED
@@ -18,10 +18,7 @@ def example_inference():
|
|
18 |
model_input_size = [1024,1024]
|
19 |
orig_im = io.imread(im_path)
|
20 |
orig_im_size = orig_im.shape[0:2]
|
21 |
-
image = preprocess_image(orig_im, model_input_size)
|
22 |
-
|
23 |
-
if torch.cuda.is_available():
|
24 |
-
image=image.cuda()
|
25 |
|
26 |
# inference
|
27 |
result=net(image)
|
|
|
18 |
model_input_size = [1024,1024]
|
19 |
orig_im = io.imread(im_path)
|
20 |
orig_im_size = orig_im.shape[0:2]
|
21 |
+
image = preprocess_image(orig_im, model_input_size).to(device)
|
|
|
|
|
|
|
22 |
|
23 |
# inference
|
24 |
result=net(image)
|