OriLib commited on
Commit
e6421bf
1 Parent(s): 398a3d9

Update example_inference.py

Browse files
Files changed (1) hide show
  1. 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)