Fixes 1adrianb#16

ghtwht · Oct 7, 2017 · 2aa7c57 · 2aa7c57
1 parent 77ec1ce
commit 2aa7c57
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/examples/detect_landmarks_in_image.py b/examples/detect_landmarks_in_image.py
@@ -38,4 +38,4 @@
 
 ax.view_init(elev=90., azim=90.)
 ax.set_xlim(ax.get_xlim()[::-1])
-plt.show()
+plt.show()
diff --git a/face_alignment/utils.py b/face_alignment/utils.py
@@ -209,8 +209,13 @@ def shuffle_lr(parts, pairs=None):
 
 
 def flip(tensor, is_label=False):
+    was_cuda = False
     if isinstance(tensor, torch.Tensor):
         tensor = tensor.numpy()
+    elif isinstance(tensor, torch.cuda.FloatTensor):
+        tensor = tensor.cpu().numpy()
+        was_cuda = True
+
     was_squeezed = False
     if tensor.ndim == 4:
         tensor = np.squeeze(tensor)
@@ -223,4 +228,7 @@ def flip(tensor, is_label=False):
         tensor = cv2.flip(tensor, 1).reshape(tensor.shape)
     if was_squeezed:
         tensor = np.expand_dims(tensor, axis=0)
-    return torch.from_numpy(tensor)
+    tensor = torch.from_numpy(tensor)
+    if was_cuda:
+       tensor = tensor.cuda()
+    return tensor