fix bug:预测图像没有在cuda上

yizt · Aug 10, 2020 · 2f62684 · 2f62684
1 parent 5b672d3
commit 2f62684
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/rest.py b/rest.py
@@ -58,14 +58,13 @@ def inference(image, h, w):
     :return: text
     """
     image = torch.FloatTensor(image)
+    image = image.to(device)
 
     if h > w:
         predict = v_net(image)[0].detach().cpu().numpy()  # [W,num_classes]
     else:
         predict = h_net(image)[0].detach().cpu().numpy()  # [W,num_classes]
 
-    image.to(device)
-
     label = np.argmax(predict[:], axis=1)
     label = [alpha[class_id] for class_id in label]
     label = [k for k, g in itertools.groupby(list(label))]