From 0d123fda29f57f787b21f2cea0e7bf39309e03b6 Mon Sep 17 00:00:00 2001 From: magehrig Date: Tue, 28 Sep 2021 10:55:56 +0200 Subject: [PATCH] coords_grid on device --- core/utils/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/utils/utils.py b/core/utils/utils.py index 5f32d28..741ccfe 100644 --- 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)