From e68c2a23e2ddc5432fb93f537056b49b869cb547 Mon Sep 17 00:00:00 2001 From: Julien RIPOCHE Date: Tue, 14 Sep 2021 11:48:26 +0200 Subject: [PATCH] Change torch.arange dtype from torch.float32 to torch.int32 in anchor_utils.py (#4395) --- torchvision/models/detection/anchor_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/models/detection/anchor_utils.py b/torchvision/models/detection/anchor_utils.py index 06ecc551442..0057da45e24 100644 --- a/torchvision/models/detection/anchor_utils.py +++ b/torchvision/models/detection/anchor_utils.py @@ -97,10 +97,10 @@ def grid_anchors(self, grid_sizes: List[List[int]], strides: List[List[Tensor]]) # For output anchor, compute [x_center, y_center, x_center, y_center] shifts_x = torch.arange( - 0, grid_width, dtype=torch.float32, device=device + 0, grid_width, dtype=torch.int32, device=device ) * stride_width shifts_y = torch.arange( - 0, grid_height, dtype=torch.float32, device=device + 0, grid_height, dtype=torch.int32, device=device ) * stride_height shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x) shift_x = shift_x.reshape(-1)