coords_grid on device
This commit is contained in:
parent
e6e53c4e23
commit
0d123fda29
@ -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)
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user