iicd/RAFT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

coords_grid on device

Browse Source
This commit is contained in:
magehrig 2021-09-28 10:55:56 +02:00
parent e6e53c4e23
commit 0d123fda29
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
core/utils/utils.py
View File

@ -71,8 +71,8 @@ def bilinear_sampler(img, coords, mode='bilinear', mask=False):
return img return img
def coords_grid(batch, ht, wd): def coords_grid(batch, ht, wd, device):
coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
coords = torch.stack(coords[::-1], dim=0).float() coords = torch.stack(coords[::-1], dim=0).float()
return coords[None].repeat(batch, 1, 1, 1) return coords[None].repeat(batch, 1, 1, 1)