coords_grid on device

hm-ysjiang · Sep 28, 2021 · 0d123fd · 0d123fd
1 parent e6e53c4
commit 0d123fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/core/utils/utils.py b/core/utils/utils.py
@@ -71,8 +71,8 @@ def bilinear_sampler(img, coords, mode='bilinear', mask=False):
     return img
 
 
-def coords_grid(batch, ht, wd):
-    coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
+def coords_grid(batch, ht, wd, device):
+    coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
     coords = torch.stack(coords[::-1], dim=0).float()
     return coords[None].repeat(batch, 1, 1, 1)