Change dtype from torch.float32 to torch.int32 in anchor_utils.py (#166)

zhiqwang · Sep 16, 2021 · b9c6720 · b9c6720
1 parent 992b8c3
commit b9c6720
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/yolort/models/anchor_utils.py b/yolort/models/anchor_utils.py
@@ -68,8 +68,8 @@ def grid_anchors(
             grid_height, grid_width = size
 
             # For output anchor, compute [x_center, y_center, x_center, y_center]
-            shifts_x = torch.arange(0, grid_width, dtype=torch.float32, device=device).to(dtype=dtype)
-            shifts_y = torch.arange(0, grid_height, dtype=torch.float32, device=device).to(dtype=dtype)
+            shifts_x = torch.arange(0, grid_width, dtype=torch.int32, device=device).to(dtype=dtype)
+            shifts_y = torch.arange(0, grid_height, dtype=torch.int32, device=device).to(dtype=dtype)
 
             shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)