[Enhance] Optimize IO for flow_warp (open-mmlab#881)

9vivian88 · May 10, 2022 · 334d003 · 334d003
1 parent fa272ac
commit 334d003
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/mmedit/models/common/flow_warp.py b/mmedit/models/common/flow_warp.py
@@ -29,8 +29,11 @@ def flow_warp(x,
                          f'flow ({flow.size()[1:3]}) are not the same.')
     _, _, h, w = x.size()
     # create mesh grid
-    grid_y, grid_x = torch.meshgrid(torch.arange(0, h), torch.arange(0, w))
-    grid = torch.stack((grid_x, grid_y), 2).type_as(x)  # (h, w, 2)
+    device = flow.device
+    grid_y, grid_x = torch.meshgrid(
+        torch.arange(0, h, device=device, dtype=x.dtype),
+        torch.arange(0, w, device=device, dtype=x.dtype))
+    grid = torch.stack((grid_x, grid_y), 2)  # h, w, 2
     grid.requires_grad = False
 
     grid_flow = grid + flow